# passerelle - uniform access to multiple data sources and services
# Copyright (C) 2018 Entr'ouvert
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import calendar
from datetime import datetime

ENCODE_TOKENS = {
    'integer': 16,
    'real': 19,
    'nil': 0,
    'true': 1,
    'false': 2,
    'emptyString': 3,
    'emptyData': 4,
    'ref': 9,
    'i1': 10,
    'u1': 11,
    'i2': 12,
    'u2': 13,
    'i4': 14,
    'uint32': 15,
    'i8': 16,
    'u8': 17,
    'float': 18,
    'double': 19,
    'decimal': 20,
    'string': 21,
    'localdate': 22,
    'gmtDate': 23,
    'color': 24,
    'data': 25,
    'naturals': 26,
    'dictionary': 30,
    'array': 31,
    'couple': 32,
    'object': 50,
}

DECODE_TOKENS = {v: k for k, v in ENCODE_TOKENS.items()}


class Couple(list):
    pass


class Uint32(int):
    pass


class MSTEDecoder:
    def __init__(self, data):
        self._idx = 4
        self._keys = []
        self._refs = []
        self._data = data

    def _next_token(self):
        self._idx += 1
        return self._data[self._idx]

    def _parse_array(self):
        count = self._next_token()
        res = []
        self._refs.append(res)
        while count > 0:
            res.append(self._parse_item())
            count -= 1
        return res

    def _parse_couple(self):
        res = Couple()
        self._refs.append(res)
        for _ in range(2):
            res.append(self._parse_item())
        return res

    def _parse_decimal(self):
        res = float(self._next_token())
        self._refs.append(res)
        return res

    def _parse_dictionary(self):
        count = self._next_token()
        res = {}
        self._refs.append(res)
        while count > 0:
            key = self._keys[self._next_token()]
            res[key] = self._parse_item()
            count -= 1
        return res

    def _parse_emptyString(self):
        return ''

    def _parse_gmtDate(self):
        return self._parse_localdate()

    def _parse_item(self):
        token = self._next_token()
        _type = DECODE_TOKENS[token]
        return getattr(self, "_parse_%s" % _type)()

    def _parse_localdate(self):
        timestamp = self._next_token()
        res = datetime.fromtimestamp(timestamp)
        self._refs.append(res)
        return res

    def _parse_nil(self):
        return None

    def _parse_ref(self):
        pos = self._next_token()
        return self._refs[pos]

    def _parse_string(self):
        res = self._next_token()
        self._refs.append(res)
        return res

    def _parse_true(self):
        return True

    def _parse_false(self):
        return False

    def _parse_uint32(self):
        return int(self._next_token())

    def decode(self):
        num_keys = self._data[self._idx]
        while num_keys > 0:
            self._keys.append(self._next_token())
            num_keys -= 1
        return self._parse_item()


class ObjectStore(list):
    def add(self, obj):
        """Add object in the store
        and return its reference
        """
        ref = self.getref(obj)
        if ref is None:
            self.append(obj)
            ref = len(self) - 1
        return ref

    def getref(self, obj):
        """Return the reference of obj,
        None if the object is not in the store
        """
        try:
            return self.index(obj)
        except ValueError:
            return None


class MSTEEncoder:
    def __init__(self, data):
        self._data = data
        self._stream = []
        self._refs_store = ObjectStore()
        self._keys_store = ObjectStore()

    def _push_token_type(self, token_type):
        self._stream.append(ENCODE_TOKENS[token_type])

    def _push(self, item):
        self._stream.append(item)

    def _encode_array(self, obj):
        self._refs_store.add(obj)
        self._push_token_type('array')
        self._push(len(obj))
        for elem in obj:
            self._encode_obj(elem)

    def _encode_boolean(self, obj):
        if obj:
            self._push_token_type('true')
        else:
            self._push_token_type('false')

    def _encode_couple(self, obj):
        self._refs_store.add(obj)
        self._push_token_type('couple')
        for elem in obj:
            self._encode_obj(elem)

    def _encode_decimal(self, obj):
        self._refs_store.add(obj)
        self._push_token_type('decimal')
        self._push(int(obj))

    def _encode_dictionary(self, obj):
        self._refs_store.add(obj)
        self._push_token_type('dictionary')
        self._push(len(obj))
        for key, value in obj.items():
            key_ref = self._keys_store.add(key)
            self._push(key_ref)
            self._encode_obj(value)

    def _encode_localdate(self, obj):
        # obj must be utc
        self._refs_store.add(obj)
        self._push_token_type('localdate')
        self._push(int(calendar.timegm(obj.timetuple())))

    def _encode_obj(self, obj):
        ref = self._refs_store.getref(obj)
        if ref is not None:
            self._push_token_type('ref')
            self._push(ref)
        elif isinstance(obj, str):
            self._encode_string(obj)
        elif obj is None:
            self._encode_nil()
        elif isinstance(obj, Couple):
            self._encode_couple(obj)
        elif isinstance(obj, bool):
            self._encode_boolean(obj)
        elif isinstance(obj, list):
            self._encode_array(obj)
        elif isinstance(obj, dict):
            self._encode_dictionary(obj)
        elif isinstance(obj, float):
            self._encode_decimal(obj)
        elif isinstance(obj, datetime):
            self._encode_localdate(obj)
        elif isinstance(obj, Uint32):
            self._encode_uint32(obj)
        else:
            raise TypeError("%s encoding not supported" % type(obj))

    def _encode_nil(self):
        self._push_token_type('nil')

    def _encode_string(self, _str):
        if _str:
            self._push_token_type('string')
            self._push(_str)
            self._refs_store.add(_str)
        else:
            self._push_token_type('emptyString')

    def _encode_uint32(self, obj):
        self._push_token_type('uint32')
        self._push(obj)

    def encode(self):
        res = ["MSTE0102"]
        self._encode_obj(self._data)
        nb_token = 5 + len(self._keys_store) + len(self._stream)
        res = ["MSTE0102", nb_token, 'CRC00000000', 0, len(self._keys_store)] + self._keys_store
        res.extend(self._stream)
        return res


def decode(data):
    return MSTEDecoder(data).decode()


def encode(data):
    return MSTEEncoder(data).encode()
