import hashlib
from bitarray import bitarray
import typing

from .. import Slice, CellTypes
from ..address import Address
from ..cell import Cell
from ..builder import Builder
from ..deserialize import NullCell

from .utils import serialize_dict
from .parse import parse_hashmap

Key = typing.Union[int, str, bytes, Address]
Value = typing.Union


class DictError(Exception):
    pass


class HashMap:

    def __init__(self,
                 key_size: int,
                 key_serializer: typing.Optional[typing.Callable] = None,
                 value_serializer: typing.Optional[typing.Callable] = None,
                 map_: typing.Optional[dict] = None,
                 ):

        self.size = key_size
        if map_ is None:
            map_ = {}
        self.map = map_
        self.key_serializer: typing.Optional[typing.Callable] = key_serializer
        self.value_serializer: typing.Optional[typing.Callable] = value_serializer

    def set_int_key(self, int_key: int, value):
        if int_key.bit_length() > self.size:
            raise DictError('Key sizes must be the same.')
        self.map[int_key] = value
        return self

    def set(self, key: Key, value, hash_key=False):
        """
        :param key: key
        :param value: value
        :param hash_key: need sha256 hash key? Usually used for tokens onchain metadata.
        :return: self
        Usage examples:
            dict = HashMap(256, value_serializer=lambda src, dest: dest.store_string(src))
            dict.set('name', 'pytoniq', hash_key=True).set('description', 'the best lib', hash_key=True)

            dict = HashMap(267).with_coins_values()
            dict.set(key=Address('EQBvW8Z5huBkMJYdnfAEM5JqTNkuWX3diqYENkWsIL0XggGG'), value=15)\
                    .set(key=Address('EQCD39VS5jcptHL8vMjEXrzGaRcCVYto7HUn4bpAOg8xqB2N'), value=10)
        """
        if self.key_serializer:
            key = self.key_serializer(key)
            if not isinstance(key, int):
                raise DictError('key serializer should return int')
            return self.set_int_key(key, value)

        if hash_key:
            if not isinstance(key, str):
                raise DictError('hash_key=True requires a string key')
            key = hashlib.sha256(key.encode()).digest()
        if isinstance(key, bytes):
            key = int.from_bytes(key, 'big', signed=False)
        elif isinstance(key, str):
            key = int(key, 2)
        elif isinstance(key, Address):
            key = Builder().store_address(key).end_cell().begin_parse().load_uint(267)
        if isinstance(key, int):
            return self.set_int_key(key, value)
        else:
            raise DictError('unknown key type')

    def with_address_values(self):
        self.value_serializer = lambda src, dest: dest.store_address(src)
        return self

    def with_uint_values(self, length: int):
        self.value_serializer = lambda src, dest: dest.store_uint(src, length)
        return self

    def with_int_values(self, length: int):
        self.value_serializer = lambda src, dest: dest.store_int(src, length)
        return self

    def with_coins_values(self):
        self.value_serializer = lambda src, dest: dest.store_coins(src)
        return self

    def serialize(self) -> typing.Optional[Cell]:
        if not self.value_serializer:
            self.value_serializer = lambda src, dest: dest.store_cell(src)
        if len(self.map):
            return serialize_dict(self.map, self.size, self.value_serializer).end_cell()
        else:
            return None

    @classmethod
    def from_cell(cls, dict_cell: Cell, key_length: int) -> "HashMap":
        dict_result = parse_hashmap(dict_cell.begin_parse(), key_length)
        map_ = {int(i, 2): j for i, j in dict_result.items()}
        result = cls(key_length)
        result.map = map_
        return result

    @staticmethod
    def parse(dict_cell: Slice,  # Slice
              key_length: int,  # bitlength of keys
              key_deserializer: typing.Optional[typing.Callable] = None,  # func to deserialize keys
              value_deserializer: typing.Optional[typing.Callable] = None  # func to deserialize values
              ) -> typing.Optional[dict]:

        if dict_cell.type_ != CellTypes.ordinary:
            return None

        if not key_deserializer:
            key_deserializer = lambda i: int(i, 2)  # by default key_deserializer just converts bitstring to int
        dict_result = parse_hashmap(dict_cell, key_length)
        if value_deserializer:
            result = {key_deserializer(i): value_deserializer(j) for i, j in dict_result.items()}
        else:
            result = {key_deserializer(i): j for i, j in
                      dict_result.items()}  # if you do not provide value_deserializer the values are Slices
        return result

    #
    # @classmethod
    # def with_address_keys(cls):
    #     return cls(267,
    #                key_serializer=lambda src: Builder().store_address(src).end_cell().begin_parse().load_bits(267),
    #                value_serializer=None)
    #
    # def with_address_values(self):
    #     self.value_serializer = lambda src, dest: dest.store_address(src)
    #     return self
    #
    # @classmethod
    # def with_uint_keys(cls, bit_length: int):
    #     return cls(
    #         bit_length,
    #         key_serializer=lambda k: Builder().store_uint(k).end_cell().begin_parse().load_bits(267),
    #
    #     )