from lib.fileutils import *
from dataclasses import dataclass, field
from typing import List, Union


@dataclass
class Vec3Indexed:
    x: float
    y: float
    z: float
    index: int

    @classmethod
    def create(cls, b, offset=None):
        old = b.tell()
        if offset:
            b.seek(offset)
        x = read_float_and_print("Vec3->x", b)
        y = read_float_and_print("Vec3->y", b)
        z = read_float_and_print("Vec3->z", b)
        index = read_int_and_print("Vec3->index", b)
        if offset:
            b.seek(old)
        return cls(x=x, y=y, z=z, index=index)


@dataclass
class Vec3:
    x: float
    y: float
    z: float

    @classmethod
    def create(cls, b, offset=None):
        old = b.tell()
        if offset:
            b.seek(offset)
        x = read_float_and_print("Vec3->x", b)
        y = read_float_and_print("Vec3->y", b)
        z = read_float_and_print("Vec3->z", b)
        if offset:
            b.seek(old)
        return cls(x=x, y=y, z=z)


@dataclass
class AnmObj:
    custom: int
    skin: int = field(repr=False)
    zero_cluster: int
    total_number_vertices: int
    num_clusters: int
    clusters: List['AnmCluster']

    @classmethod
    def create(cls, b, offset=None):
        old = b.tell()
        if offset:
            b.seek(offset)
        custom = read_int_and_print("ANM_Obj->Custom", b)
        skin = read_int_and_print("ANM_Obj->Skin", b)
        zero = read_int_and_print("ANM_Obj->ZeroCluster", b)
        total_vertices = read_int_and_print("ANM_Obj->TotalNumberOfVertices", b)
        num_clusters = read_int_and_print("ANM_Obj->NumberOfClusters", b)
        clusters_offset = read_int_and_print("ANM_Obj->Clusters", b)
        clusters = []
        _old1 = b.tell()
        b.seek(clusters_offset)
        for _ in range(num_clusters):
            cluster = AnmCluster.create(b)
            clusters.append(cluster)

        b.seek(_old1)
        if offset:
            b.seek(old)
        r = cls(custom=custom, skin=skin, zero_cluster=zero, total_number_vertices=total_vertices, num_clusters=num_clusters, clusters=clusters)
        for cluster in clusters:
            cluster.obj = r
        return r


@dataclass
class AnmCluster:
    custom: int = field(repr=False)
    obj: Union[AnmObj, None]
    handle: int
    bounding_box: List[Vec3] = field(repr=False)  # size 2
    num_vertices: int
    vertices: List[Vec3Indexed] = field(repr=False)
    bone_name: str = ""

    @classmethod
    def create(cls, b, offset=None):
        old = b.tell()
        if offset:
            b.seek(offset)
        custom = read_int_and_print("ANM_Cluster->Custom", b)
        obj_offset = read_int_and_print("ANM_Cluster->Obj", b)
        if obj_offset == 0:  # ugly hack, dunno why we need that
            _ = read_int_and_print("ANM_Cluster->Custom", b)
            _ = read_int_and_print("ANM_Cluster->Obj", b)
        obj = None  # is set by parent Obj
        handle = read_int_and_print("ANM_Cluster->Handle", b)
        bbox1 = Vec3.create(b)
        bbox2 = Vec3.create(b)
        bounding_box = [bbox1, bbox2]
        num_vertices = read_int_and_print("ANM_Cluster->num_vertices", b)
        vertices = []
        for _ in range(num_vertices):
            v = Vec3Indexed.create(b)
            vertices.append(v)

        if offset:
            b.seek(old)
        return cls(custom=custom, obj=obj, handle=handle, bounding_box=bounding_box, num_vertices=num_vertices, vertices=vertices)


class NamDictionnary:
    @staticmethod
    def create(b, cls, depth=1, offset=None):
        r = {}

        old = b.tell()
        if offset:
            b.seek(offset)

        read_int_and_print("NAM_Dictionnary->pMemGroup", b)
        namemapper = read_int_and_print("NAM_Dictionnary->pNameMapper", b)
        flags = read_int_and_print("NAM_Dictionnary->Flags", b)
        print("Flags: NAM_FlagHasHandles: %s, NAM_FlagCaseInsensitive: %s" % (flags & NAM_FlagHasHandles, flags & NAM_FlagCaseInsensitive))
        read_short_and_print("NAM_Dictionnary->NameHandle", b)

        b.seek(namemapper)
        read_int_and_print("pMemGroup", b)
        read_int_and_print("MemBaseSize", b)
        read_int_and_print("MemBlocks", b)
        read_int_and_print("MemSize", b)
        read_int_and_print("MemFree", b)
        read_int_and_print("pMem", b)
        read_int_and_print("MaxMappings", b)
        totalmappings = read_int_and_print("TotalMappings", b)
        read_int_and_print("TotalKeys", b)
        p_key_contexts = read_int_and_print("pKeyContexts", b)

        print("MAP_Image->ANM_Skin->NAM_Dictionnary->MAP_Mapper->MAP_KeyContext")
        b.seek(p_key_contexts)
        print("\tMAP_Image->ANM_Skin->NAM_Dictionnary->MAP_Mapper->MAP_KeyContext->MAP_BinSearchContext")
        _ = read_int_and_print("\tSize", b)
        read_int_and_print("\tpCallback", b)
        read_int_and_print("\tIndex", b)
        read_int_and_print("\tTest", b)
        p_lut = read_int_and_print("pLUT", b)
        read_int_and_print("Reserve", b)
        read_int_and_print("pData", b)

        for i in range(totalmappings):
            b.seek(p_lut)
            pindx = read_int(b)
            for _ in range(i):
                pindx = read_int(b)
            print("LUT INDEX: %s" % pindx)
            b.seek(pindx)
            if not flags & NAM_FlagHasHandles:
                read_short_and_print("pPtr", b)
                read_short_and_print("pPtr2", b)
                val = read_int_and_print("Value", b)
                name = read_string_until_none_and_print("name", b)
                if depth > 1:
                    item = NamDictionnary.create(b, cls, depth-1, val)
                else:
                    item = cls.create(b, val)
                r[name] = item
            else:
                refcnt = read_short_and_print("RefCnt", b)
                handle = read_short_and_print("Handle", b)
                read_int_and_print("Test", b)
                read_int_and_print("Test", b)
                bone_name = read_string_until_none_and_print("BoneName", b)
                r[str(i)] = {"RefCnt": refcnt, "Handle": handle, "BoneName": bone_name}

        if offset:
            b.seek(old)

        return r


NAM_FlagHasHandles = (1 << 0)
NAM_FlagCaseInsensitive = (1 << 1)


def read_int_and_print(n, b):
    a = read_int(b)
    print("%s=%s" % (n, a))
    return a


def read_float_and_print(n, b):
    a = read_float(b)
    print("%s=%s" % (n, a))
    return a


def read_byte_and_print(n, b):
    a = read_byte(b)
    print("%s=%s" % (n, a))
    return a


def read_short_and_print(n, b):
    a = read_short(b)
    print("%s=%s" % (n, a))
    return a


def read_string_until_none_and_print(n, b):
    s = read_string_until_none(b)
    print("%s=%s" % (n, s))
    return s


def link(nodedict, defaultobjtree):
    for obj_name, obj in defaultobjtree.items():
        for cluster in obj.clusters:
            cluster_handle = cluster.handle
            k = str(cluster_handle-1)
            cluster.bone_name = nodedict[k]["BoneName"]


class AnmSkin:
    def __init__(self):
        self.custom = 0
        self.file_version = 0
        self.image = 0
        self.memgroup = 0
        self.node_dictionary = 0
        self.tree_dictionary = 0
        self.default_obj_dictionnary = 0
        self.total_number_trees = 0
        self.total_number_objs = 0
        self.total_number_clusters = 0
        self.name = 0

    @classmethod
    def parse(cls, b):
        _ = read_int_and_print("size", b)
        memgroupoffset = read_int_and_print("memgroupoffset", b)

        b.seek(memgroupoffset)

        print("MAP_MemGroup")
        print("\tMAP_Mem")
        read_int_and_print("\tNode", b)
        read_int_and_print("\tSize", b)
        read_int_and_print("\tFlags", b)

        read_int_and_print("pMemGroup", b)
        read_int_and_print("MemList", b)
        read_int_and_print("TotalSize", b)
        read_int_and_print("TotalBlocks", b)
        base = read_int_and_print("pBase", b)

        b.seek(base)
        print("=========")
        print("ANM_Skin")
        read_int_and_print("pCustom", b)
        read_int_and_print("FileVersion", b)
        read_int_and_print("pImage", b)
        read_int_and_print("pMemGroup", b)
        node = read_int_and_print("pNodeDictionary", b)
        tree = read_int_and_print("pTreeDictionary", b)
        defaultobj = read_int_and_print("pDefaultObjDictionary", b)
        read_int_and_print("TotalNumberOfTrees", b)
        read_int_and_print("TotalNumberOfObjs", b)
        read_int_and_print("TotalNumberOfClusters", b)
        read_string_until_none_and_print("Name", b)
        print("=========")

        nodedict = NamDictionnary.create(b, AnmObj, offset=node)
        treedict = NamDictionnary.create(b, AnmObj, depth=2, offset=tree)
        defaultobjdict = NamDictionnary.create(b, AnmObj, offset=defaultobj)

        link(nodedict, defaultobjdict)

        print("nodedict_len:%s" % len(nodedict))
        print("treedict_len:%s" % len(treedict))
        print(defaultobjdict)
        print("defaultobjtreedict_len:%s" % len(defaultobjdict))


def main():
    with open("/home/tasty/Projects/gck-map-extract-objects/anm_skn/mc_l0.skn", "rb") as fp:
        AnmSkin.parse(fp)

    with open("/home/tasty/Projects/gck-map-extract-objects/anm_skn/rp_l0.skn", "rb") as fp:
        AnmSkin.parse(fp)

    with open("/home/tasty/Projects/gck-map-extract-objects/anm_skn/kb_l0.skn", "rb") as fp:
        AnmSkin.parse(fp)


if __name__ == "__main__":
    main()