import os
from typing import List, Union
from lib.packet import Packet
from lib.fileutils import *
from lib.game import Vec3, VecRGB, cross
import argparse
import sys

GBS_VERSION = 0xaa0100be
GBSFlagNormals = 0x0001
GBSFlagUVs = 0x0002
GBSFlagRGBs = 0x0004
GBSFlagCalcNormals = 0x0008
GBSFlagMaxLit = (1 << 31)


def resize(l: List, t, num):
    l.clear()
    for i in range(num):
        l.append(t())


class FileMaxObj:
    def __init__(self):
        self.vstart: int = 0
        self.vcount: int = 0
        self.nstart: int = 0
        self.ncount: int = 0
        self.noffset: int = 0


class MaxObj(FileMaxObj):
    def __init__(self):
        super().__init__()
        self.fstart: int = 0
        self.fcount: int = 0
        self.sostart: int = 0
        self.socount: int = 0


class SubObject:
    def __init__(self):
        self.objname: str = ""
        self.maxobjindex: int = 0
        self.ntris: int = 0  # count of tridata (including preceding 'count' short)
        self.totaltris: int = 0
        self.tridata: List[int] = []  # unsigned short
        self.verticeref_start: int = 0
        self.verticeref_count: int = 0
        self.texname: str = ""
        self.bumptexture: str = ""
        self.falloff: float = 0
        self.blend: float = 0
        self.flags: int = 0
        self.emissive: int = 0
        self.ambient: int = 0
        self.diffuse: int = 0
        self.specular: int = 0
        self.power: int = 0


class UV:
    def __init__(self, u: float = 0, v: float = 0):
        self.u: float = u
        self.v: float = v


class OBJMaterial:
    def __init__(self):
        self.name = ""
        self.texture = ""


class OBJObject:
    def __init__(self):
        self.faces: List[OBJFace] = []
        self.name = "root"
        self.material: Union[None, OBJMaterial] = None


class OBJFace:
    def __init__(self):
        self.index_vertices: List[int] = []
        self.index_uvs: List[int] = []
        self.index_normals: List[int] = []


def obj_read_materials(matlib_file) -> List[OBJMaterial]:
    materials = []
    curr_mat = None
    with open(matlib_file, "r") as fp:
        while line := fp.readline():
            line = line.strip()
            arr = line.split(" ")
            if arr[0] == "newmtl":
                mat = OBJMaterial()
                materials.append(mat)
                mat.name = arr[1].rstrip()
                print("NEW MAT: %s" % mat.name)
                curr_mat = mat
            if arr[0] == "map_Ka" or arr[0] == "map_Kd":
                matname_without_ext = "".join(arr[1:]).split("/")[-1]
                matname_without_ext = "".join(matname_without_ext.split(".")[0:-1])
                curr_mat.texture = matname_without_ext
                print("MAT %s HAS TEXTURE %s" % (curr_mat.name, curr_mat.texture))
    return materials


class GbsData:
    def __init__(self):
        self.name = ""
        self.optionsflags: int = 0
        self.nndefs: int = 0
        self.num_normals: int = 0
        self.normals: List[int] = []  # word
        self.indexed_normals: List[int] = []  # unsigned short
        self.num_vertices: int = 0
        self.vertices: List[Vec3] = []
        self.nsobjs: int = 0
        self.nmobjs: int = 0
        self.indexed_vertices: List[int] = []  # unsigned short
        self.vertrgb: List[VecRGB] = []
        self.nverts: int = 0
        self.vertuv: List[UV] = []
        self.MaxObjs: List[MaxObj] = []
        self.SubObjs: List[SubObject] = []

    def read(self, file):
        debug = True
        self.name = os.path.basename(file)
        with open(file, "rb") as fp:
            version_header = read_int(fp)
            if version_header != GBS_VERSION:
                raise Exception("File does not appear to be a GBS file.")

            self.optionsflags = read_int(fp)
            self.num_vertices = read_int(fp)
            resize(self.vertices, Vec3, self.num_vertices)
            for i in range(self.num_vertices):
                self.vertices[i].x = read_float(fp)
                self.vertices[i].y = read_float(fp)
                self.vertices[i].z = read_float(fp)

            if self.optionsflags & GBSFlagNormals:
                self.nndefs = read_int(fp)
                self.num_normals = read_int(fp)
                resize(self.normals, int, self.num_normals)
                for i in range(self.num_normals):
                    self.normals[i] = read_short(fp)

            self.nverts = read_int(fp)
            if debug: print("NVERTS=%s, NUMVERTICES=%s, options=%s" % (self.nverts, self.num_vertices, self.optionsflags))
            resize(self.indexed_vertices, int, self.nverts)
            for i in range(self.nverts):
                self.indexed_vertices[i] = read_short(fp)
            if debug: print(self.indexed_vertices)

            if self.optionsflags & GBSFlagNormals:
                resize(self.indexed_normals, int, self.nverts)
                for i in range(self.nverts):
                    self.indexed_normals[i] = read_short(fp)

            if self.optionsflags & GBSFlagUVs:
                resize(self.vertuv, UV, self.nverts)
                if debug: print("Read %s UV" % self.nverts)
                for i in range(self.nverts):
                    self.vertuv[i].u = read_float(fp)
                    self.vertuv[i].v = read_float(fp) * -1

            if self.optionsflags & GBSFlagRGBs:
                resize(self.vertrgb, VecRGB, self.nverts)
                if debug: print("Read %s RGB" % self.nverts)
                for i in range(self.nverts):
                    self.vertrgb[i].r = read_byte(fp)
                    self.vertrgb[i].g = read_byte(fp)
                    self.vertrgb[i].b = read_byte(fp)

            # Get number of objects
            self.nmobjs = read_int(fp)
            if debug: print("NMOBJS = %s" % self.nmobjs)
            resize(self.MaxObjs, MaxObj, self.nmobjs)
            for maxobj in self.MaxObjs:
                print("Processing new maxobj")
                fileMaxObj = FileMaxObj()
                fileMaxObj.vstart = read_int(fp)
                fileMaxObj.vcount = read_int(fp)
                fileMaxObj.nstart = read_int(fp)
                fileMaxObj.ncount = read_int(fp)
                fileMaxObj.noffset = read_int(fp)
                print("vstart=%s vcount=%s nstart=%s ncount=%s noffset=%s" % (fileMaxObj.vstart, fileMaxObj.vcount, fileMaxObj.nstart, fileMaxObj.ncount, fileMaxObj.noffset))
                maxobj.vstart = fileMaxObj.vstart
                maxobj.vcount = fileMaxObj.vcount
                maxobj.nstart = fileMaxObj.nstart
                maxobj.ncount = fileMaxObj.ncount
                maxobj.noffset = fileMaxObj.noffset
                maxobj.fstart = 0
                maxobj.fcount = 0
                maxobj.sostart = 0
                maxobj.socount = 0

            # verify max obj
            nmcount = 0
            for maxobj in self.MaxObjs:
                nmcount += maxobj.vcount
            assert nmcount == self.num_vertices

            self.nsobjs = read_int(fp)
            print("Subobject count: %s" % self.nsobjs)
            resize(self.SubObjs, SubObject, self.nsobjs)
            num_faces = 0
            len_tridata = 0
            for ns in range(self.nsobjs):
                if debug: print("Processing subobj %s" % ns)
                object = self.SubObjs[ns]
                object.objname = read_string(fp, 32)
                object.maxobjindex = read_int(fp)
                object.totaltris = read_int(fp)
                if debug: print("read %s totaltris" % object.totaltris)
                num_faces += object.totaltris
                object.ntris = read_int(fp)

                assert ((object.ntris - 1) / 3 == object.totaltris)

                resize(object.tridata, int, object.ntris + 1)
                if debug: print("read gbs: totaltris: %s, tridata is %s long, ntris: %s" % (object.totaltris, len(object.tridata), object.ntris))
                len_tridata += object.ntris + 1
                for i in range(object.ntris):
                    object.tridata[i] = read_short(fp)

                # if debug: print("tridata: %s" % object.tridata)

                object.verticeref_start = read_int(fp)
                object.verticeref_count = read_int(fp)
                if self.optionsflags & GBSFlagUVs:
                    object.texname = read_string(fp, 32)
                    object.bumptexture = read_string(fp, 32)

                object.falloff = read_float(fp)
                if self.optionsflags & GBSFlagRGBs:
                    object.blend = read_float(fp)

                object.flags = read_int(fp)
                object.emissive = read_int(fp)
                object.ambient = read_int(fp)
                object.diffuse = read_int(fp)
                object.specular = read_int(fp)
                object.power = read_float(fp)
                if debug: print("%s (%s) falloff: %s, blend: %s, flags:%s, emissive: %s, ambiant: %s, diffuse: %s, specular: %s, power: %s" % (object.objname, object.texname, object.falloff, object.blend, object.flags, object.emissive, object.ambient, object.diffuse, object.specular, object.power))

                maxobj = self.MaxObjs[object.maxobjindex]
                maxobj.fcount += object.totaltris
                if not maxobj.socount:
                    maxobj.socount = ns

                maxobj.socount += 1
            if debug: print("read num_faces: %s, len_tridata: %s" % (num_faces, len_tridata))


    @staticmethod
    def evaluate_tridata(tridata, tri_idx, count):
        if count == 0:
            count = tridata[0]
            if count == 0:
                return False
            tri_idx = 0
        v1 = tridata[tri_idx + 1]
        v2 = tridata[tri_idx + 2]
        v3 = tridata[tri_idx + 3]
        tri_idx += 3

        count -= 1
        if count < 0:
            count = 0xffff  # max unsigned short
        if count == 0:
            tridata = tridata[tridata[0] * 3 + 1:]
        return tridata, tri_idx, count, v1, v2, v3

    def generate_normals(self):
        normals: List[Vec3] = []
        resize(normals, Vec3, self.num_vertices)
        for subobj_i in range(len(self.SubObjs)):
            subobj = self.SubObjs[subobj_i]
            tridata = subobj.tridata
            values = self.evaluate_tridata(tridata, -1, 0)
            while values:
                tridata, tri_idx, count, v1, v2, v3 = values
                try:
                    p = cross(self.vertices[self.indexed_vertices[v2]] - self.vertices[self.indexed_vertices[v1]], self.vertices[self.indexed_vertices[v3]] - self.vertices[self.indexed_vertices[v1]])
                except Exception as e:
                    raise e
                normals[self.indexed_vertices[v1]] += p
                normals[self.indexed_vertices[v2]] += p
                normals[self.indexed_vertices[v3]] += p
                values = self.evaluate_tridata(tridata, tri_idx, count)

        for i in range(len(normals)):
            normals[i] = normals[i].normalize()
        return normals

    @staticmethod
    def from_obj(obj_file):
        materials = []
        vertices: List[Vec3] = []
        uvs: List[UV] = []
        normals: List[Vec3] = []
        objects: List[OBJObject] = []
        root_obj = OBJObject()
        objects.append(root_obj)

        last_material = None
        current_object = root_obj

        uv_ind = {}

        with open(obj_file, "r") as obj_fp:
            while line := obj_fp.readline():
                arr = line.split(" ")
                if arr[0] == "v":
                    v = Vec3(float(arr[1]), float(arr[2]), float(arr[3]))
                    vertices.append(v)
                if arr[0] == "vt":
                    v = UV(float(arr[1]), float(arr[2]))
                    uvs.append(v)
                if arr[0] == "vn":
                    v = Vec3(float(arr[1]), float(arr[2]), float(arr[3]))
                    normals.append(v)
                if arr[0] == "f":
                    f = OBJFace()
                    f1_s = arr[1].split("/")
                    f2_s = arr[2].split("/")
                    f3_s = arr[3].split("/")
                    if len(f1_s) == len(f2_s) == len(f3_s) >= 1:
                        # face only has vertex index data
                        v1_index = int(f1_s[0])
                        v2_index = int(f2_s[0])
                        v3_index = int(f3_s[0])
                        f.index_vertices.append(v1_index)
                        f.index_vertices.append(v2_index)
                        f.index_vertices.append(v3_index)
                    if len(f1_s) == len(f2_s) == len(f3_s) >= 2:
                        # face has vertex index and uv
                        v1_uv_index = int(f1_s[1])
                        v2_uv_index = int(f2_s[1])
                        v3_uv_index = int(f3_s[1])
                        uv_ind[v1_uv_index] = v1_index
                        uv_ind[v2_uv_index] = v2_index
                        uv_ind[v3_uv_index] = v3_index
                        f.index_uvs.append(v1_uv_index)
                        f.index_uvs.append(v2_uv_index)
                        f.index_uvs.append(v3_uv_index)
                    if len(f1_s) == len(f2_s) == len(f3_s) >= 3:
                        # face has vertex index and uv and normal
                        v1_normal_index = int(f1_s[2])
                        v2_normal_index = int(f2_s[2])
                        v3_normal_index = int(f3_s[2])
                        f.index_normals.append(v1_normal_index)
                        f.index_normals.append(v2_normal_index)
                        f.index_normals.append(v3_normal_index)
                    current_object.faces.append(f)
                if arr[0] == "o":
                    obj_name = arr[1].rstrip()
                    o = OBJObject()
                    o.name = obj_name
                    o.material = last_material
                    objects.append(o)
                    if len(current_object.faces) == 0:
                        objects.remove(current_object)
                    current_object = o
                if arr[0] == "usemtl":
                    mtl_name = arr[1].rstrip()
                    print("mtl_name: %s" % mtl_name)
                    mtl = [mat for mat in materials if mat.name == mtl_name][0]
                    print("FOUND MTL: %s for name %s" % (mtl.texture, mtl_name))
                    current_object.material = mtl
                    last_material = mtl
                if arr[0] == "mtllib":
                    matlib_file = arr[1].rstrip()
                    print("Reading mtllib %s" % matlib_file)
                    materials = obj_read_materials("%s/%s" % (os.path.dirname(obj_file), matlib_file))

            num_faces = sum([len(o.faces) for o in objects])
            print("%s vertices, %s uvs, %s normals, %s objects, %s materials, %s faces" % (len(vertices), len(uvs), len(normals), len(objects), len(materials), num_faces))

            try:
                assert(len(vertices) < 32000)
            except AssertionError:
                print("Your model has %s vertices. Giants only supports a maximum of 32k vertices" % len(vertices))
                sys.exit(1)

            try:
                assert(len(uv_ind) < 0xffff)
            except AssertionError:
                print("Your model has %s UV. Giants only supports a maximum of 65k UV per model" % len(uv_ind))
                sys.exit(1)

            len_vertices = len(vertices)
            len_normals = len(normals)

            data = Packet()
            data.put_ulong(GBS_VERSION)
            options = 0
            if len(uvs) > 0:
                options |= GBSFlagUVs
            options |= GBSFlagCalcNormals
            data.put_long(options)
            data.put_long(len_vertices)
            for v in vertices:
                data.put_float(v.x)
                data.put_float(v.y)
                data.put_float(v.z)

            if options & GBSFlagNormals:
                for v in normals:
                    data.put_float(v.x)
                    data.put_float(v.y)
                    data.put_float(v.z)

                data.put_ulong(len_normals)
                data.put_ulong(len_normals)
                for i in range(len_normals):
                    data.put_short(len_vertices + i)

            assert len_vertices <= len(uvs)

            indices = []
            for uv_in in sorted(uv_ind.keys()):
                indices.append(uv_ind[uv_in]-1)

            nverts = max(len_vertices, len(uvs))
            assert nverts == len(indices)

            data.put_ulong(nverts)
            for i in indices:
                # indexed vert/UV
                data.put_short(i)

            if options & GBSFlagUVs:
                for v in uvs:
                    data.put_float(v.u)
                    data.put_float(v.v * -1)

            # max objects
            data.put_long(1)  # 1 big object
            data.put_long(0)  # vstart
            data.put_long(len_vertices)  # vcount
            data.put_long(0)  # nstart
            data.put_long(0)  # ncount
            data.put_long(0)  # noffset ???

            # start write subobjects
            data.put_long(len(objects))
            for obj in objects:
                print("Writing object %s" % obj.name)
                data.put_string_size(obj.name, 32)
                data.put_long(0)  # max obj index
                data.put_long(len(obj.faces))  # totaltris
                print("wrote %s totaltris" % len(obj.faces))
                data.put_long(3 * len(obj.faces) + 1)  # ntris
                print("wrote %s ntris" % (3 * len(obj.faces) + 1))

                data.put_short(len(obj.faces))
                for face in obj.faces:  # obj.faces.length == ntris
                    data.put_short(face.index_uvs[0] - 1)
                    data.put_short(face.index_uvs[1] - 1)
                    data.put_short(face.index_uvs[2] - 1)

                data.put_long(0)  # verticeref_start
                data.put_long(nverts)  # verticeref_count
                if options & GBSFlagUVs:
                    print("Writing texture to GBS: %s" % obj.material.texture)
                    data.put_string_size(obj.material.texture, 32)  # texture
                    data.put_string_size(obj.material.texture, 32)  # bump

                data.put_float(0)  # falloff
                if options & GBSFlagRGBs:
                    data.put_float(0)  # blend

                data.put_long(0x40000000)  # flags
                data.put_long(0)  # emissive
                data.put_long(0)  # ambient
                data.put_long(0)  # diffuse
                data.put_long(0)  # specular
                data.put_float(0)  # power

            # write gbs file
            output = "%s/%s.gbs" % (os.path.dirname(os.path.abspath(obj_file)), os.path.basename(obj_file))
            with open(output, "wb") as gbs:
                gbs.write(data.getvalue())

    def __str__(self):
        return "[name: %s, nverts:%s, num_vertices: %s, vertuv: %s, normals: %s, maxobjs: %s, subobjs:%s]" % (self.name, self.nverts, self.num_vertices, len(self.vertuv), self.num_normals, self.nmobjs, self.nsobjs)


def convert_obj(path):
    GbsData.from_obj(path)
    output = "%s/%s.gbs" % (os.path.dirname(os.path.abspath(path)), os.path.basename(path))
    print("Done! Output: %s" % output)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("path")
    args = parser.parse_args()
    if os.path.exists(args.path):
        convert_obj(args.path)
    else:
        print("ERROR: file %s does not exist" % args.path)