from __future__ import annotations

import re
from sys import version_info as py_version_info
from typing import TYPE_CHECKING, Any, Optional, Union

from attrs import define

from .. import classes
from ..streams.EndianBinaryReader import EndianBinaryReader
from ..streams.EndianBinaryWriter import EndianBinaryWriter
from .TypeTreeNode import TypeTreeNode

Object = classes.Object
UnknownObject = classes.UnknownObject
PPtr = classes.PPtr

if TYPE_CHECKING:
    from ..files.SerializedFile import SerializedFile


try:
    from ..UnityPyBoost import read_typetree as read_typetree_boost
except ImportError:
    read_typetree_boost = None

kAlignBytes = 0x4000

FUNCTION_READ_MAP = {
    "SInt8": EndianBinaryReader.read_byte,
    "UInt8": EndianBinaryReader.read_u_byte,
    "char": EndianBinaryReader.read_u_byte,
    "short": EndianBinaryReader.read_short,
    "SInt16": EndianBinaryReader.read_short,
    "unsigned short": EndianBinaryReader.read_u_short,
    "UInt16": EndianBinaryReader.read_u_short,
    "int": EndianBinaryReader.read_int,
    "SInt32": EndianBinaryReader.read_int,
    "unsigned int": EndianBinaryReader.read_u_int,
    "UInt32": EndianBinaryReader.read_u_int,
    "Type*": EndianBinaryReader.read_u_int,
    "long long": EndianBinaryReader.read_long,
    "SInt64": EndianBinaryReader.read_long,
    "unsigned long long": EndianBinaryReader.read_u_long,
    "UInt64": EndianBinaryReader.read_u_long,
    "FileSize": EndianBinaryReader.read_u_long,
    "float": EndianBinaryReader.read_float,
    "double": EndianBinaryReader.read_double,
    "bool": EndianBinaryReader.read_boolean,
    "string": EndianBinaryReader.read_aligned_string,
    "TypelessData": EndianBinaryReader.read_byte_array,
}

FUNCTION_READ_MAP_ARRAY = {
    "SInt8": EndianBinaryReader.read_byte_array,
    "UInt8": EndianBinaryReader.read_u_byte_array,
    "char": EndianBinaryReader.read_u_byte_array,
    "short": EndianBinaryReader.read_short_array,
    "SInt16": EndianBinaryReader.read_short_array,
    "unsigned short": EndianBinaryReader.read_u_short_array,
    "UInt16": EndianBinaryReader.read_u_short_array,
    "int": EndianBinaryReader.read_int_array,
    "SInt32": EndianBinaryReader.read_int_array,
    "unsigned int": EndianBinaryReader.read_u_int_array,
    "UInt32": EndianBinaryReader.read_u_int_array,
    "Type*": EndianBinaryReader.read_u_int_array,
    "long long": EndianBinaryReader.read_long_array,
    "SInt64": EndianBinaryReader.read_long_array,
    "unsigned long long": EndianBinaryReader.read_u_long_array,
    "UInt64": EndianBinaryReader.read_u_long_array,
    "FileSize": EndianBinaryReader.read_u_long_array,
    "float": EndianBinaryReader.read_float_array,
    "double": EndianBinaryReader.read_double_array,
    "bool": EndianBinaryReader.read_boolean_array,
}


@define(slots=True)
class TypeTreeConfig:
    as_dict: bool
    assetsfile: "Optional[SerializedFile]" = None
    has_registry: bool = False

    def copy(self) -> TypeTreeConfig:
        return TypeTreeConfig(self.as_dict, self.assetsfile, self.has_registry)


def get_ref_type_node(ref_object: dict, assetfile: SerializedFile) -> Optional[TypeTreeNode]:
    typ = ref_object["type"]
    if isinstance(typ, dict):
        cls = typ["class"]
        ns = typ["ns"]
        asm = typ["asm"]
    else:
        cls = getattr(typ, "class")
        ns = typ.ns
        asm = typ.asm

    if not assetfile or not assetfile.ref_types:
        raise ValueError("SerializedFile has no ref_types")

    if cls == "":
        return None

    for ref_type in assetfile.ref_types:
        if cls == ref_type.m_ClassName and ns == ref_type.m_NameSpace and asm == ref_type.m_AssemblyName:
            return ref_type.node
    else:
        raise ValueError(f"Referenced type not found: {cls} {ns} {asm}")


if py_version_info >= (3, 14):
    from annotationlib import get_annotations as annotationlib_get_annotations

    def get_annotation_keys(clz) -> set[str]:
        return set(annotationlib_get_annotations(clz).keys())
else:

    def get_annotation_keys(clz) -> set[str]:
        return set(clz.__annotations__)


def read_typetree(
    root_node: TypeTreeNode,
    reader: EndianBinaryReader,
    as_dict: bool = True,
    byte_size: Optional[int] = None,
    check_read: bool = True,
    assetsfile: Optional[SerializedFile] = None,
) -> Union[dict[str, Any], Object]:
    """Reads the typetree of the object contained in the reader via the node list.

    Parameters
    ----------
    nodes : list
        List of nodes/nodes
    reader : EndianBinaryReader
        Reader of the object to be parsed

    Returns
    -------
    dict | objects.Object
        The parsed typtree
    """
    bytes_read: int
    if byte_size and read_typetree_boost:
        data = reader.read_bytes(byte_size)
        obj, bytes_read = read_typetree_boost(data, root_node, reader.endian, as_dict, assetsfile, classes)
    else:
        pos = reader.Position
        config = TypeTreeConfig(as_dict, assetsfile, False)
        obj = read_value(root_node, reader, config)
        bytes_read = reader.Position - pos

    if check_read and bytes_read != byte_size:
        raise ValueError(f"Expected to read {byte_size} bytes, but only read {bytes_read} bytes")

    return obj


def write_typetree(
    value: Union[dict[str, Any], Object],
    root_node: TypeTreeNode,
    writer: EndianBinaryWriter,
    assetsfile: Optional[SerializedFile] = None,
) -> None:
    """Writes the typetree of the object contained in the reader via the node list.

    Parameters
    ----------
    value: dict | objects.Object
        The object to be written
    nodes : list
        List of nodes/nodes
    writer : EndianBinaryWriter
        Writer of the object to be parsed
    """
    config = TypeTreeConfig(isinstance(value, dict), assetsfile, False)
    return write_value(value, root_node, writer, config)


def read_value(
    node: TypeTreeNode,
    reader: EndianBinaryReader,
    config: TypeTreeConfig,
) -> Any:
    # print(reader.Position, node.m_Name, node.m_Type, node.m_MetaFlag)
    align = metaflag_is_aligned(node.m_MetaFlag)

    func = FUNCTION_READ_MAP.get(node.m_Type)
    if func:
        value = func(reader)
    elif node.m_Type == "pair":
        first = read_value(node.m_Children[0], reader, config)
        second = read_value(node.m_Children[1], reader, config)
        value = (first, second)
    elif node.m_Type == "ReferencedObject":
        value = {}
        for child in node.m_Children:
            if child.m_Type == "ReferencedObjectData":
                ref_type_nodes = get_ref_type_node(value, config.assetsfile)
                if ref_type_nodes is None:
                    continue
                value[child.m_Name] = read_value(ref_type_nodes, reader, config)
            else:
                value[child.m_Name] = read_value(child, reader, config)
    # Vector
    elif node.m_Children and node.m_Children[0].m_Type == "Array":
        if metaflag_is_aligned(node.m_Children[0].m_MetaFlag):
            align = True

        # size = read_value(node.m_Children[0].m_Children[0], reader, as_dict)
        size = reader.read_int()
        if size < 0:
            raise ValueError("Negative length read from TypeTree")
        subtype = node.m_Children[0].m_Children[1]
        if metaflag_is_aligned(subtype.m_MetaFlag):
            value = read_value_array(subtype, reader, config, size)
        else:
            value = [read_value(subtype, reader, config) for _ in range(size)]

    else:  # Class
        value = {}
        for child in node.m_Children:
            if child.m_Type == "ManagedReferencesRegistry":
                if config.has_registry:
                    continue
                else:
                    config = config.copy()
                    config.has_registry = True
            value[child.m_Name if config.as_dict else child._clean_name] = read_value(child, reader, config)

        if not config.as_dict:
            if node.m_Type.startswith("PPtr<"):
                value = PPtr[Any](
                    assetsfile=config.assetsfile,
                    m_FileID=value["m_FileID"],
                    m_PathID=value["m_PathID"],
                )
            else:
                clz = getattr(classes, node.m_Type, UnknownObject)
                try:
                    value = clz(**value)
                except TypeError:
                    keys = set(value.keys())
                    annotation_keys = get_annotation_keys(clz)
                    missing_keys = annotation_keys - keys
                    if clz is UnknownObject or missing_keys:
                        value = UnknownObject(node, **value)
                    else:
                        extra_keys = keys - annotation_keys
                        if extra_keys:
                            instance = clz(**{key: value[key] for key in annotation_keys})
                            for key in extra_keys:
                                setattr(instance, key, value[key])
                            value = instance
                        else:
                            value = UnknownObject(**value)

    if align:
        reader.align_stream()

    return value


def read_value_array(
    node: TypeTreeNode,
    reader: EndianBinaryReader,
    config: TypeTreeConfig,
    size: int,
) -> Any:
    align = metaflag_is_aligned(node.m_MetaFlag)

    func = FUNCTION_READ_MAP_ARRAY.get(node.m_Type)
    if func:
        value = func(reader, size)
    elif node.m_Type == "string":
        value = [reader.read_aligned_string() for _ in range(size)]
    elif node.m_Type == "TypelessData":
        value = [reader.read_bytes() for _ in range(size)]
    elif node.m_Type == "pair":
        key_node = node.m_Children[0]
        value_node = node.m_Children[1]

        key_func = FUNCTION_READ_MAP.get(
            key_node.m_Type,
            lambda reader: read_value(key_node, reader, config),
        )
        value_func = FUNCTION_READ_MAP.get(
            value_node.m_Type,
            lambda reader: read_value(value_node, reader, config),
        )
        value = [(key_func(reader), value_func(reader)) for _ in range(size)]
    elif node.m_Type == "ReferencedObject":
        value = [None] * size
        for i in range(size):
            item = {}
            for child in node.m_Children:
                if child.m_Type == "ReferencedObjectData":
                    ref_type_nodes = get_ref_type_node(item, config.assetsfile)
                    item[child.m_Name] = read_value(ref_type_nodes, reader, config)
                else:
                    item[child.m_Name] = read_value(child, reader, config)
            value[i] = item
    # Vector
    elif node.m_Children and node.m_Children[0].m_Type == "Array":
        if metaflag_is_aligned(node.m_Children[0].m_MetaFlag):
            align = True
        subtype = node.m_Children[0].m_Children[1]
        if metaflag_is_aligned(subtype.m_MetaFlag):
            value = [read_value_array(subtype, reader, config, reader.read_int()) for _ in range(size)]
        else:
            value = [[read_value(subtype, reader, config) for _ in range(reader.read_int())] for _ in range(size)]
    else:  # Class
        if config.as_dict:
            value = [
                {child.m_Name: read_value(child, reader, config) for child in node.m_Children} for _ in range(size)
            ]
        elif node.m_Type.startswith("PPtr<"):
            value = [
                PPtr[Any](
                    assetsfile=config.assetsfile,
                    **{child.m_Name: read_value(child, reader, config) for child in node.m_Children},
                )
                for _ in range(size)
            ]
        else:
            clz = getattr(
                classes,
                node.m_Type,
                UnknownObject,
            )
            keys = set(child._clean_name for child in node.m_Children)
            annotation_keys = get_annotation_keys(clz)
            missing_keys = annotation_keys - keys
            extra_keys = keys - annotation_keys
            if missing_keys or clz is UnknownObject:
                value = [
                    UnknownObject(
                        node,
                        **{child._clean_name: read_value(child, reader, config) for child in node.m_Children},
                    )
                    for _ in range(size)
                ]
            elif extra_keys:
                value = [None] * size
                for i in range(size):
                    value_i_d = {child._clean_name: read_value(child, reader, config) for child in node.m_Children}
                    value_i = clz(**{key: value for key, value in value_i_d.items() if key in annotation_keys})
                    for key in extra_keys:
                        setattr(value_i, key, value_i_d[key])
                    value[i] = value_i
            else:
                value = [
                    clz(**{child._clean_name: read_value(child, reader, config) for child in node.m_Children})
                    for _ in range(size)
                ]

    if align:
        reader.align_stream()
    return value


def metaflag_is_aligned(meta_flag: int | None) -> bool:
    return ((meta_flag or 0) & kAlignBytes) != 0


def clean_name(name: str) -> str:
    if name.startswith("(int&)"):
        name = name[6:]
    if name.endswith("?"):
        name = name[:-1]
    name = re.sub(r"[ \.:\-\[\]]", "_", name)
    if name in ["pass", "from"]:
        name += "_"
    if name[0].isdigit():
        name = f"x{name}"
    return name


FUNCTION_WRITE_MAP = {
    "SInt8": EndianBinaryWriter.write_byte,
    "UInt8": EndianBinaryWriter.write_u_byte,
    "char": EndianBinaryWriter.write_u_byte,
    "short": EndianBinaryWriter.write_short,
    "SInt16": EndianBinaryWriter.write_short,
    "unsigned short": EndianBinaryWriter.write_u_short,
    "UInt16": EndianBinaryWriter.write_u_short,
    "int": EndianBinaryWriter.write_int,
    "SInt32": EndianBinaryWriter.write_int,
    "unsigned int": EndianBinaryWriter.write_u_int,
    "UInt32": EndianBinaryWriter.write_u_int,
    "Type*": EndianBinaryWriter.write_u_int,
    "long long": EndianBinaryWriter.write_long,
    "SInt64": EndianBinaryWriter.write_long,
    "unsigned long long": EndianBinaryWriter.write_u_long,
    "UInt64": EndianBinaryWriter.write_u_long,
    "FileSize": EndianBinaryWriter.write_u_long,
    "float": EndianBinaryWriter.write_float,
    "double": EndianBinaryWriter.write_double,
    "bool": EndianBinaryWriter.write_boolean,
    "string": EndianBinaryWriter.write_aligned_string,
    "TypelessData": EndianBinaryWriter.write_byte_array,
}


def write_value(
    value: Union[dict[str, Any], Object],
    node: TypeTreeNode,
    writer: EndianBinaryWriter,
    config: TypeTreeConfig,
) -> None:
    # print(reader.Position, node.m_Name, node.m_Type, node.m_MetaFlag)
    align = metaflag_is_aligned(node.m_MetaFlag)

    func = FUNCTION_WRITE_MAP.get(node.m_Type)
    if func:
        value = func(writer, value)
    elif node.m_Type == "pair":
        write_value(value[0], node.m_Children[0], writer, config)
        write_value(value[1], node.m_Children[1], writer, config)
    elif node.m_Type == "ReferencedObject":
        for child in node.m_Children:
            if child.m_Type == "ReferencedObjectData":
                ref_type_nodes = get_ref_type_node(value, config.assetsfile)
                write_value(value[child.m_Name], ref_type_nodes, writer, config)
            else:
                write_value(value[child.m_Name], child, writer, config)
    elif node.m_Children and node.m_Children[0].m_Type == "Array":
        if metaflag_is_aligned(node.m_Children[0].m_MetaFlag):
            align = True

        writer.write_int(len(value))
        subtype = node.m_Children[0].m_Children[1]
        [write_value(sub_value, subtype, writer, config) for sub_value in value]

    else:  # Class
        if isinstance(value, dict):
            for child in node.m_Children:
                if child.m_Type == "ManagedReferencesRegistry":
                    if config.has_registry:
                        continue
                    else:
                        config = config.copy()
                        config.has_registry = True
                write_value(value[child.m_Name], child, writer, config)
        else:
            for child in node.m_Children:
                if child.m_Type == "ManagedReferencesRegistry":
                    if config.has_registry:
                        continue
                    else:
                        config = config.copy()
                        config.has_registry = True
                write_value(getattr(value, child._clean_name), child, writer, config)

    if align:
        writer.align_stream()
