import gc
import struct
import typing
from typing import Tuple, List

from pypinch.consts import NUMBER_BASE, ObjType, POSITIVE_INT_FLAG, NULL_FLAG, BYTES_FLAG, \
    LIST_FLAG, \
    DICT_FLAG, STR_KEY_DICT_FLAG, FLOAT_FLAG, STR_FLAG, NEGATIVE_INT_FLAG, \
    EMPTY_LIST_FLAG, EMPTY_DICT_FLAG, CONSISTENT_TYPE_LIST_FLAG, INT_FLAG, BOOL_FLAG, POINTER_FLAG, \
    ByteLike, HEADER, BIG_ENDIAN_DOUBLE_FORMAT, NUMBER_OF_BITS_IN_BYTE, \
    LEFTMOST_BIT_MASK, BYTES_IN_DOUBLE, FIRST_FLAGS_LIST, AMOUNT_OF_USED_FLAGS, \
    INVALID_UTF_8_START_BYTE_COMPACT_ASCII, ASCII_STR_FLAG, LIST_OF_STRUCTURED_DICTS_FLAG

from pypinch.exceptions import DeserializationError
from pypinch.deserialize.settings import Settings
from pypinch.deserialize.utils import decode_number


def load_bytes(
    buffer: ByteLike,
    *,
    use_tuples: bool = False,
    stop_gc: bool = False,
    ignore_extra_data: bool = False
) -> ObjType:

    try:
        if stop_gc:
            gc.freeze()

        settings = Settings(
            use_tuples=use_tuples,  # TODO
            pointers=[],
        )
        res, pointer = deserialize_object(buffer, len(HEADER), settings)
        if not ignore_extra_data and pointer != len(buffer):
            raise DeserializationError("Extra data")
        return res
    except DeserializationError:
        raise
    except MemoryError:
        raise
    except Exception as e:
        raise DeserializationError() from e
    finally:
        if stop_gc:
            gc.unfreeze()


def deserialize_object(buffer: bytes, pointer: int, settings: Settings) -> (ObjType, int):
    flag = buffer[pointer]
    pointer += 1
    if flag < len(FIRST_FLAGS_LIST):
        return FIRST_FLAGS_LIST[flag], pointer
    elif flag == POSITIVE_INT_FLAG:
        return decode_number(buffer, pointer)
    elif flag == STR_KEY_DICT_FLAG:
        length, pointer = decode_number(buffer, pointer)
        res_dict = {}
        for i in range(length):
            if buffer[pointer] == NUMBER_BASE - 1:
                position, pointer = decode_number(buffer, pointer + 1)
                k = settings.pointers[position]
            else:
                k, pointer = deserialize_str(buffer, pointer, settings, base=NUMBER_BASE - 1)
            v, pointer = deserialize_object(buffer, pointer, settings)
            res_dict[k] = v
        return res_dict, pointer
    elif flag == ASCII_STR_FLAG:
        return deserialize_str(buffer, pointer, settings)
    elif flag == STR_FLAG:
        return deserialize_str(buffer, pointer, settings)
    elif flag == DICT_FLAG:
        length, pointer = decode_number(buffer, pointer)
        res_dict = {}
        for _ in range(length):
            if buffer[pointer] == STR_FLAG:
                # fast path
                k, pointer = deserialize_str(buffer, pointer + 1, settings)
            else:
                k, pointer = deserialize_object(buffer, pointer, settings)
            v, pointer = deserialize_object(buffer, pointer, settings)
            res_dict[k] = v
        return res_dict, pointer
    elif flag == EMPTY_DICT_FLAG:
        return {}, pointer
    elif flag == LIST_FLAG:
        length, pointer = decode_number(buffer, pointer)
        res_list = [None] * length
        for i in range(length):
            res_list[i], pointer = deserialize_object(buffer, pointer, settings)
        return res_list, pointer
    elif flag == EMPTY_LIST_FLAG:
        return (tuple() if settings.use_tuples else []), pointer
    elif flag == CONSISTENT_TYPE_LIST_FLAG:
        typ_flag = buffer[pointer]
        length, pointer = decode_number(buffer, pointer + 1)
        if typ_flag == NULL_FLAG:
            return ((None,) if settings.use_tuples else [None]) * length, pointer
        elif typ_flag == INT_FLAG:
            res_list = typing.cast(List[int], [None] * length)
            for i in range(length):
                if buffer[pointer] == NUMBER_BASE - 1:
                    num, pointer = decode_number(buffer, pointer + 1, base=NUMBER_BASE - 1)
                    res_list[i] = -num
                else:
                    num, pointer = decode_number(buffer, pointer, base=NUMBER_BASE - 1)
                    res_list[i] = num
            return res_list, pointer
        elif typ_flag == BOOL_FLAG:
            res_list = typing.cast(List[bool], [None] * length)
            # same as: math.ceil(length / NUMBER_OF_BITS_IN_BYTE)
            # the `>> 3` is like dividing by 8 (8 is `1000` in binary)
            # the + 7 is like rounding up
            length_in_bytes = (length + 7) >> 3
            try:
                for i in range(length_in_bytes):
                    byte = buffer[pointer + i]
                    for j in range(NUMBER_OF_BITS_IN_BYTE):
                        res_list[i * NUMBER_OF_BITS_IN_BYTE + j] = (byte & LEFTMOST_BIT_MASK) == LEFTMOST_BIT_MASK
                        byte <<= 1
            except IndexError:
                pass
            return res_list, pointer + length_in_bytes
        elif typ_flag == BYTES_FLAG:
            res_list = typing.cast(List[bytes], [None] * length)
            for i in range(length):
                bytes_length, pointer = decode_number(buffer, pointer)
                res_list[i] = bytes(buffer[pointer:pointer + bytes_length])
                pointer += bytes_length
            return res_list, pointer
        elif typ_flag == STR_FLAG:
            res_list = typing.cast(List[str], [None] * length)
            for i in range(length):
                res_list[i], pointer = deserialize_str(buffer, pointer, settings)
            return res_list, pointer
        elif typ_flag == FLOAT_FLAG:
            res_list = [
                struct.unpack(BIG_ENDIAN_DOUBLE_FORMAT, buffer[pos:pos + BYTES_IN_DOUBLE])[0]
                for pos in range(pointer, pointer + BYTES_IN_DOUBLE * length, BYTES_IN_DOUBLE)
            ]
            pointer += BYTES_IN_DOUBLE * length
            return res_list, pointer
        else:
            raise DeserializationError(f"Unexpected type flag: {typ_flag}")
    elif flag == NEGATIVE_INT_FLAG:
        num, pointer = decode_number(buffer, pointer)
        return -num, pointer
    elif flag == FLOAT_FLAG:
        num = struct.unpack(BIG_ENDIAN_DOUBLE_FORMAT, buffer[pointer:pointer + BYTES_IN_DOUBLE])[0]
        return num, pointer + BYTES_IN_DOUBLE
    elif flag == BYTES_FLAG:
        length, pointer = decode_number(buffer, pointer)
        return bytes(buffer[pointer:pointer + length]), pointer + length
    elif flag == POINTER_FLAG:
        position, pointer = decode_number(buffer, pointer)
        return settings.pointers[position], pointer
    elif flag == LIST_OF_STRUCTURED_DICTS_FLAG:
        list_length, pointer = decode_number(buffer, pointer)
        dict_length, pointer = decode_number(buffer, pointer)
        res_list = typing.cast(List[typing.Dict], [None] * list_length)
        # first dict:
        first_dict = {}
        keys = typing.cast(List[str], [None] * dict_length)
        for i in range(dict_length):
            k, pointer = deserialize_object(buffer, pointer, settings)
            v, pointer = deserialize_object(buffer, pointer, settings)
            first_dict[k] = v
            keys[i] = k
        res_list[0] = first_dict
        # rest of the dicts:
        for list_idx in range(1, list_length):
            dct = {}
            for i, key in enumerate(keys):
                dct[key], pointer = deserialize_object(buffer, pointer, settings)
            res_list[list_idx] = dct
        return res_list, pointer
    elif flag < AMOUNT_OF_USED_FLAGS:
        raise DeserializationError("unexpected flag")
    else:
        return flag - AMOUNT_OF_USED_FLAGS, pointer


def deserialize_str(buffer: bytes, pointer: int, settings: Settings, base: int = NUMBER_BASE) -> Tuple[str, int]:
    length, pointer = decode_number(buffer, pointer, base=base)
    string = buffer[
        #          Skip 1 char if buffer starts with INVALID_UTF_8_START_BYTE_COMPACT_ASCII
        pointer + (buffer[pointer] == INVALID_UTF_8_START_BYTE_COMPACT_ASCII)
        :pointer + length
    ].decode()
    settings.pointers.append(string)
    return string, pointer + length
