from __future__ import annotations

import re
from struct import Struct
from threading import Lock
from typing import (
    TYPE_CHECKING,
    Dict,
    Iterator,
    List,
    Optional,
    Tuple,
    Union,
    cast,
)

from attrs import define, field

from ..streams.EndianBinaryReader import EndianBinaryReader
from ..streams.EndianBinaryWriter import EndianBinaryWriter

if TYPE_CHECKING:
    from .Tpk import UnityVersion

try:
    from ..UnityPyBoost import TypeTreeNode as TypeTreeNodeC  # type: ignore
except ImportError:

    @define(slots=True)
    class TypeTreeNodeC:
        m_Level: int
        m_Type: str
        m_Name: str
        m_ByteSize: int
        m_Version: int
        m_Children: List[TypeTreeNode] = field(factory=list)
        m_TypeFlags: Optional[int] = None
        m_VariableCount: Optional[int] = None
        m_Index: Optional[int] = None
        m_MetaFlag: Optional[int] = None
        m_RefTypeHash: Optional[int] = None
        _clean_name: str = field(init=False)

        def __attrs_post_init__(self):
            self._clean_name = clean_name(self.m_Name)

        def __repr__(self):
            return f"TypeTreeNode(m_Level={self.m_Level}, m_Type='{self.m_Type}', \
                m_Name='{self.m_Name}', m_MetaFlag={self.m_MetaFlag})"


TYPETREENODE_KEYS = [
    "m_Level",
    "m_Type",
    "m_Name",
    "m_ByteSize",
    "m_Version",
    "m_Children",
    "m_TypeFlags",
    "m_VariableCount",
    "m_Index",
    "m_MetaFlag",
    "m_RefTypeHash",
]

SYSTEM_GLOBAL_LOCK = Lock()
NAME_PEEK_NODE_CACHE: dict[Tuple[str, str, int], Union[Tuple[TypeTreeNode, str], None]] = {}


class TypeTreeNode(TypeTreeNodeC):
    def traverse(self) -> Iterator[TypeTreeNode]:
        stack: list[TypeTreeNode] = [self]
        while stack:
            node = stack.pop()
            yield node
            stack.extend(reversed(node.m_Children))

    @classmethod
    def parse(cls, reader: EndianBinaryReader, version: int) -> TypeTreeNode:
        # stack approach is way faster than recursion
        # using a fake root node to avoid special case for root node
        dummy_node = cls(-1, "", "", 0, 0, [])
        dummy_root = cls(-1, "", "", 0, 0, [dummy_node])

        stack: List[Tuple[TypeTreeNode, int]] = [(dummy_root, 1)]
        while stack:
            parent, count = stack[-1]
            if count == 1:
                stack.pop()
            else:
                stack[-1] = (parent, count - 1)

            node = cls(
                m_Level=parent.m_Level + 1,
                m_Type=reader.read_string_to_null(),
                m_Name=reader.read_string_to_null(),
                m_ByteSize=reader.read_int(),
                m_VariableCount=reader.read_int() if version == 2 else None,
                m_Index=reader.read_int() if version != 3 else None,
                m_TypeFlags=reader.read_int(),
                m_Version=reader.read_int(),
                m_MetaFlag=reader.read_int() if version != 3 else None,
            )
            parent.m_Children[-count] = node
            children_count = reader.read_int()
            if children_count > 0:
                node.m_Children = [dummy_node] * children_count
                stack.append((node, children_count))
        return dummy_root.m_Children[0]

    @classmethod
    def parse_blob(cls, reader: EndianBinaryReader, version: int) -> TypeTreeNode:
        node_count = reader.read_int()
        stringbuffer_size = reader.read_int()

        node_struct, keys = _get_blob_node_struct(reader.endian, version)
        struct_data = reader.read(node_struct.size * node_count)
        stringbuffer_reader = EndianBinaryReader(reader.read(stringbuffer_size), reader.endian)

        CommonString = get_common_strings()

        def read_string(reader: EndianBinaryReader, value: int) -> str:
            is_offset = (value & 0x80000000) == 0
            if is_offset:
                reader.Position = value
                return reader.read_string_to_null()

            offset = value & 0x7FFFFFFF
            return CommonString.get(offset, str(offset))

        fake_root: TypeTreeNode = cls(-1, "", "", 0, 0, [])
        stack: List[TypeTreeNode] = [fake_root]
        parent = fake_root
        prev = fake_root

        for raw_node in node_struct.iter_unpack(struct_data):
            node = cls(
                **dict(zip(keys[:3], raw_node[:3])),
                **dict(zip(keys[5:], raw_node[5:])),
                m_Type=read_string(stringbuffer_reader, raw_node[3]),
                m_Name=read_string(stringbuffer_reader, raw_node[4]),
            )

            if node.m_Level > prev.m_Level:
                stack.append(parent)
                parent = prev
            elif node.m_Level < prev.m_Level:
                while node.m_Level <= parent.m_Level:
                    parent = stack.pop()

            parent.m_Children.append(node)
            prev = node

        return fake_root.m_Children[0]

    @classmethod
    def from_list(cls, nodes: Union[List[Dict[str, Union[str, int]]], List[TypeTreeNode]]) -> TypeTreeNode:
        fake_root: TypeTreeNode = cls(-1, "", "", 0, 0, [])
        stack: List[TypeTreeNode] = [fake_root]
        parent = fake_root
        prev = fake_root

        # check if the nodes contain all required fields
        if isinstance(nodes[0], dict):
            if "m_Level" not in nodes[0] or "m_Type" not in nodes[0] or "m_Name" not in nodes[0]:
                raise ValueError("Nodes must contain at least m_Level, m_Type and m_Name")
            patch_dict = {}
            if "m_ByteSize" not in nodes[0]:
                patch_dict["m_ByteSize"] = 0
            if "m_Version" not in nodes[0]:
                patch_dict["m_Version"] = 0
            nodes = [cls(**node, **patch_dict) for node in nodes]  # type: ignore

        if TYPE_CHECKING:
            nodes = cast(List[TypeTreeNode], nodes)

        for node in nodes:
            if node.m_Level > prev.m_Level:
                stack.append(parent)
                parent = prev
            elif node.m_Level < prev.m_Level:
                while node.m_Level <= parent.m_Level:
                    parent = stack.pop()

            parent.m_Children.append(node)
            prev = node

        return fake_root.m_Children[0]

    def get_name_peek_node(self) -> Union[Tuple[TypeTreeNode, str], None]:
        global SYSTEM_GLOBAL_LOCK
        with SYSTEM_GLOBAL_LOCK:
            key = (self.m_Name, self.m_Type, self.m_Version)
            if key in NAME_PEEK_NODE_CACHE:
                return NAME_PEEK_NODE_CACHE[key]

            result: Union[Tuple[TypeTreeNode, str], None] = None
            for i, child in enumerate(self.m_Children):
                if child.m_Name in ("m_Name", "name"):
                    peek_node = TypeTreeNode(
                        self.m_Level,
                        self.m_Type,
                        self.m_Name,
                        self.m_ByteSize,
                        self.m_Version,
                        self.m_Children[: i + 1],
                    )
                    result = peek_node, child.m_Name
                    break
            NAME_PEEK_NODE_CACHE[key] = result
            return result

    def dump(self, writer: EndianBinaryWriter, version: int):
        stack: list[TypeTreeNode] = [self]
        while stack:
            node = stack.pop()

            writer.write_string_to_null(self.m_Type)
            writer.write_string_to_null(self.m_Name)
            writer.write_int(self.m_ByteSize)
            if version == 2:
                assert self.m_VariableCount is not None
                writer.write_int(self.m_VariableCount)
            if version != 3:
                assert self.m_Index is not None
                writer.write_int(self.m_Index)
            writer.write_int(self.m_TypeFlags or 0)
            writer.write_int(self.m_Version)
            if version != 3:
                assert self.m_MetaFlag is not None
                writer.write_int(self.m_MetaFlag)

            writer.write_int(len(self.m_Children))

            stack.extend(reversed(node.m_Children))

    def dump_blob(self, writer: EndianBinaryWriter, version: int):
        node_writer = EndianBinaryWriter(endian=writer.endian)
        string_writer = EndianBinaryWriter()

        # string buffer setup
        CommonStringOffsetMap = {string: offset for offset, string in get_common_strings().items()}

        string_offsets: dict[str, int] = {}

        def write_string(string: str) -> int:
            offset = string_offsets.get(string)
            if offset is None:
                common_offset = CommonStringOffsetMap.get(string)
                if common_offset:
                    offset = common_offset | 0x80000000
                else:
                    offset = string_writer.Position
                    string_writer.write_string_to_null(string)
                string_offsets[string] = offset
            return offset

        # node buffer setup
        node_struct, keys = _get_blob_node_struct(writer.endian, version)

        def write_node(node: TypeTreeNode):
            node_writer.write(
                node_struct.pack(
                    *[getattr(node, key) for key in keys[:3]],
                    write_string(node.m_Type),
                    write_string(node.m_Name),
                    *[getattr(node, key) for key in keys[5:]],
                )
            )

        # write nodes
        node_count = len([write_node(node) for node in self.traverse()])

        # write blob
        writer.write_int(node_count)
        writer.write_int(string_writer.Position)
        writer.write(node_writer.bytes)
        writer.write(string_writer.bytes)

    def dump_structure(self, indent: str = "  ") -> str:
        # dump structure similar to https://github.com/AssetRipper/TypeTreeDumps/blob/main/StructsDump
        sb = [
            f"{indent}{self.m_Type} {self.m_Name} // ByteSize{{{self.m_ByteSize:X}}}, Index{{{self.m_Index}}}, \
                Version{{{self.m_Version}}}, TypeFlags{{{self.m_TypeFlags}}}, MetaFlag{{{self.m_MetaFlag}}}"
        ]
        for child in self.m_Children:
            sb.append(child.dump_structure(indent + "  "))
        return "\n".join(sb)

    def to_dict(self) -> dict:
        return {
            key: value for key, value in ((key, getattr(self, key)) for key in TYPETREENODE_KEYS) if value is not None
        }

    def to_dict_list(self) -> List[dict]:
        return [
            self.to_dict(),
            *(item for child in self.m_Children for item in child.to_dict_list()),
        ]

    def __eq__(self, other: TypeTreeNode) -> bool:  # type: ignore
        return self.to_dict() == other.to_dict() and self.m_Children == other.m_Children


COMMONSTRING_CACHE: Dict[Optional[UnityVersion], Dict[int, str]] = {}


def get_common_strings(version: Optional[UnityVersion] = None) -> Dict[int, str]:
    if version in COMMONSTRING_CACHE:
        return COMMONSTRING_CACHE[version]

    from .Tpk import TPKTYPETREE

    tree = TPKTYPETREE
    common_string = tree.CommonString
    strings = common_string.GetStrings(tree.StringBuffer)
    if version:
        count = common_string.GetCount(version)
        strings = strings[:count]

    ret: Dict[int, str] = {}
    offset = 0
    for string in strings:
        ret[offset] = string
        offset += len(string) + 1

    COMMONSTRING_CACHE[version] = ret
    return ret


def _get_blob_node_struct(endian: str, version: int) -> tuple[Struct, list[str]]:
    struct_type = f"{endian}hBBIIiii"
    keys = [
        "m_Version",
        "m_Level",
        "m_TypeFlags",
        "m_TypeStrOffset",
        "m_NameStrOffset",
        "m_ByteSize",
        "m_Index",
        "m_MetaFlag",
    ]
    if version >= 19:
        struct_type += "Q"
        keys.append("m_RefTypeHash")

    return Struct(struct_type), keys


def clean_name(name: str) -> str:
    # keep in sync with TypeTreeHelper.cpp
    if len(name) == 0:
        return name
    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


__all__ = (
    "TypeTreeNode",
    "get_common_strings",
    "clean_name",
)
