"""This is an incomplete implementation of the UBJSON specification. Expected is a readable file pointer to a UBJSON file.
Things that are not implemented:
- High precision numbers
- Optimized arrays with type & count
- nested arrays, so arrays in arrays are not supported, objects in arrays are supported
"""
import struct

import numpy as np


def b(s):
    return isinstance(s, str) and s.encode('latin1') or s

NOOP = b('N')
NULL = b('Z')
FALSE = b('F')
TRUE = b('T')
INT8 = b('i')
UINT8 = b('U')
INT16 = b('I')
INT32 = b('l')
INT64 = b('L')
FLOAT32 = b('d')
FLOAT64 = b('D')
CHAR = b('C')
STRING = b('S')
HIDEF = b('H')
ARRAY_OPEN = b('[')
ARRAY_CLOSE = b(']')
OBJECT_OPEN = b('{')
OBJECT_CLOSE = b('}')

INTEGERS = [INT8, UINT8, INT16, INT32, INT64]
FLOATS = [FLOAT32, FLOAT64]

type_sizes = {
    INT8: 1,
    UINT8: 1,
    INT16: 2,
    INT32: 4,
    INT64: 8,
    FLOAT32: 4,
    FLOAT64: 8,
    # todo: maybe add high-precision number
    CHAR: 1,
}

struct_mapping = {
    INT8: 'b',
    UINT8: 'B',
    INT16: 'h',
    INT32: 'i',
    INT64: 'q',
    FLOAT32: 'f',
    FLOAT64: 'd',
    CHAR: 's',
}

numpy_type_mapping = {
    INT8: np.int8,
    UINT8: np.uint8,
    INT16: np.int16,
    INT32: np.int32,
    INT64: np.int64,
    FLOAT32: np.float32,
    FLOAT64: np.float64,
}


objects = [OBJECT_OPEN, OBJECT_CLOSE]
arrays = [ARRAY_CLOSE, ARRAY_OPEN]

# Can decode optimized array
# [[][#][i][5] // An array of 5 elements.
#     [d][29.97]
#     [d][31.13]
#     [d][67.0]
#     [d][2.113]
#     [d][23.8889]
# // No end marker since a count was specified.

# AND

# [[][$][d][#][i][5] // An array of 5 float32 elements.
#     [29.97] // Value type is known, so type markers are omitted.
#     [31.13]
#     [67.0]
#     [2.113]
#     [23.8889]
# // No end marker since a count was specified.
def _decode_array_optimized(fp):
    tag = fp.read(1)
    # optimized array with count
    if tag == b"#":
        array_length_indicator = fp.read(1)
        array_length_indicator_type = struct_mapping[array_length_indicator]
        array_length = type_sizes[array_length_indicator]
        array_length = struct.unpack(f'>{array_length_indicator_type}', fp.read(array_length))[0]

        array = []
        for _ in range(array_length):
            tag = fp.read(1)
            element = __decode_element(tag, fp)
            array.append(element)
        return array
    # optimized array with type & count
    elif tag == b"$":
        value_type_byte = fp.read(1)
        value_type_length = type_sizes[value_type_byte]
        tag = fp.read(1)
        array_type_byte = fp.read(1)
        array_type_length = type_sizes[array_type_byte]
        array_type_prefix = struct_mapping[array_type_byte]
        array_length_bytes = fp.read(array_type_length)
        array_length = struct.unpack(f'>{array_type_prefix}', array_length_bytes)[0]
        if array_length == 0:
            return list()
        buffer = fp.read(array_length * value_type_length)
        return list(struct.unpack('>' + f'{struct_mapping[value_type_byte]}' * array_length, buffer))
    else:
        raise ValueError("Expected optimized array but got received bytes of unoptimized array.")

def _decode_object(tag, fp):
    result_dict = dict()
    if tag == OBJECT_OPEN:
        key_type = None
        while key_type != OBJECT_CLOSE:
            if key_type is None:
                key_type = fp.read(1)
                # case for empty object
                if key_type == OBJECT_CLOSE:
                    return {}
            if key_type == OBJECT_OPEN:
                return _decode_object(key_type, fp)
            key, value = _decode_simple_key_value_pair(fp, key_type)
            if key == '}':
                return result_dict
            result_dict[key] = value
            key_type = fp.read(1)
    return result_dict

def __decode_element(tag, fp):
    if (element_type_length := type_sizes.get(tag)) is not None:
        value_bytes = fp.read(element_type_length)
        value_struct_prefix = struct_mapping[tag]
        return struct.unpack(f'>{value_struct_prefix}', value_bytes)[0]
    elif tag == STRING:
        string_length_type = fp.read(1)
        length = __decode_element(string_length_type, fp)
        string_bytes = fp.read(length)
        return string_bytes.decode("utf-8")
    elif tag == OBJECT_OPEN:
        return _decode_object(tag, fp)
    else:
        raise ValueError(f"Expected type size for {tag} but got {element_type_length}")

# Can decode
# [i][3][lat][d][29.976]
def _decode_simple_key_value_pair(fp, key_type):
    if key_type in type_sizes:
        length_of_key = __decode_element(key_type, fp)
        key_to_decode = fp.read(length_of_key)
        key = key_to_decode.decode("utf-8")
        value_type_byte = fp.read(1)
        if value_type_byte in type_sizes:
            value = __decode_element(value_type_byte, fp)
            return key, value
        elif value_type_byte == STRING:
            # todo: check for string, high precision number, array, object here
            # value_type_byte = __decode_element(value_type_byte, fp)
            value = __decode_element(value_type_byte, fp)
            return key, str(value)
        elif value_type_byte == OBJECT_OPEN:
            value = _decode_object(value_type_byte, fp)
            return key, value
        elif value_type_byte == OBJECT_CLOSE:
            return key, {}
        elif value_type_byte == ARRAY_OPEN:
            value = _decode_array_optimized(fp)
            return key, value
        elif value_type_byte == ARRAY_CLOSE:
            return key, []
        elif value_type_byte == b'' and key == '}':
            return key, None
        else:
            raise ValueError(f"Unmatched value type for {value_type_byte}.")
    else:
        raise ValueError(f"Expected type size for {key_type} but could not find any.")

def decode_ubjson_buffer(fp):
    fp.read(1)
    complete_dict = dict()
    key_type = fp.read(1)
    while key_type != b'' and key_type != OBJECT_CLOSE:
        key, value = _decode_simple_key_value_pair(fp, key_type)
        complete_dict[key] = value
        key_type = fp.read(1)
    return complete_dict
