import glob
import hashlib
import os
import shutil


# フォルダ一覧を取得する
def GetFolderList(RootPath):
    FileList = os.listdir(RootPath)
    DirList = [f for f in FileList if os.path.isdir(os.path.join(RootPath, f))]

    return DirList


#
class MergeTexture:

    def __init__(self, NewIndex, HEO, TextureIndexInHEO, Hash):
        self.m_NewIndex = NewIndex
        self.m_HEO = HEO
        self.m_TextureIndexInHEO = TextureIndexInHEO
        self.m_Hash = Hash

    def GetNewIndex(self):
        return self.m_NewIndex

    def GetHash(self):
        return self.m_Hash

    def GetFileName(self):
        return format(str(self.m_NewIndex), '0>3') + ".png"

    def GetSrcFileName(self):
        return format(str(self.m_TextureIndexInHEO), '0>3') + ".png"

    def CopyFile(self, SrcFolderPath, DestFolderPath):
        FolderNameList = ["tex", "tex_astc", "tex_dxt", "tex_etc"]
        ExtList = ["png", "pvr", "dds", "pvr"]

        for n in range(len(FolderNameList)):
            FolderName = FolderNameList[n]
            Ext = ExtList[n]

            SrcFolderName = SrcFolderPath + self.m_HEO.GetDir() + "/" + FolderName

            if os.path.isdir(SrcFolderName) == False:
                continue

            SrcFileName = SrcFolderName + "/" + self.GetSrcFileName()
            DestFolderName = DestFolderPath + FolderName
            DestFileName = DestFolderName + "/" + self.GetFileName()

            SrcFileName = os.path.splitext(SrcFileName)[0] + "." + Ext
            DestFileName = os.path.splitext(DestFileName)[0] + "." + Ext

            os.makedirs(DestFolderName, exist_ok=True)

            print(SrcFileName + " -> " + DestFileName)

            shutil.copyfile(SrcFileName, DestFileName)


class MergeReflectionProbe:

    def __init__(self, NewIndex, HEO, IndexInHEO, Hash):
        self.m_NewIndex = NewIndex
        self.m_HEO = HEO
        self.m_IndexInHEO = IndexInHEO
        self.m_Hash = Hash

    def GetNewIndex(self):
        return self.m_NewIndex

    def GetHash(self):
        return self.m_Hash

    def GetFolderName(self):
        return format(str(self.m_NewIndex), '0>3')

    def GetSrcFolderName(self):
        return format(str(self.m_IndexInHEO), '0>3')

    def CopyFile(self, SrcFolderPath, DestFolderPath):
        FolderNameList = ["tex_reflection_cube", "tex_reflection_cube_astc", "tex_reflection_cube_dxt",
                          "tex_reflection_cube_etc"]
        ExtList = ["png", "pvr", "dds", "pvr"]

        for n in range(len(FolderNameList)):
            FolderName = FolderNameList[n]
            Ext = ExtList[n]

            SrcFolderName = SrcFolderPath + self.m_HEO.GetDir() + "/" + FolderName + "/" + self.GetSrcFolderName()

            if os.path.isdir(SrcFolderName) == False:
                continue

            DestFolderName = DestFolderPath + FolderName + "/" + self.GetFolderName()

            os.makedirs(DestFolderName, exist_ok=True)

            for n in range(6):
                SrcFileName = SrcFolderName + "/reflection_cube_" + str(n) + "." + Ext
                DestFileName = DestFolderName + "/reflection_cube_" + str(n) + "." + Ext

                print(SrcFileName + " -> " + DestFileName)

                shutil.copyfile(SrcFileName, DestFileName)


class MergeLightMap:

    def __init__(self, NewIndex, HEO, TextureIndexInHEO, Hash):
        self.m_NewIndex = NewIndex
        self.m_HEO = HEO
        self.m_TextureIndexInHEO = TextureIndexInHEO
        self.m_Hash = Hash

    def GetNewIndex(self):
        return self.m_NewIndex

    def GetHash(self):
        return self.m_Hash

    def GetSrcFileName(self):
        return "LightMap_" + format(str(self.m_TextureIndexInHEO), '0>3') + ".png"

    def GetFileName(self):
        return "LightMap_" + format(str(self.m_NewIndex), '0>3') + ".png"

    def CopyFile(self, SrcFolderPath, DestFolderPath):
        SrcFileName = SrcFolderPath + self.m_HEO.GetDir() + "/tex/" + self.GetSrcFileName()
        DestFolderName = DestFolderPath + "tex"
        DestFileName = DestFolderName + "/" + self.GetFileName()

        os.makedirs(DestFolderName, exist_ok=True)

        print(SrcFileName + " -> " + DestFileName)

        shutil.copyfile(SrcFileName, DestFileName)


# HEOTexture
class HEOTexture:

    def __init__(self, FileNameOffset, FileName):
        self.m_FileNameOffset = FileNameOffset
        self.m_FileName = FileName
        self.m_Hash = ""
        self.m_Replaced = False

    def MakeHash(self, RootPath):
        FileName = RootPath + "tex/" + self.m_FileName

        with open(FileName, 'rb') as f:
            data = f.read()

        self.m_Hash = hashlib.md5(data).hexdigest()

        print(FileName + ", " + self.m_Hash)

    def GetHash(self):
        return self.m_Hash

    def Replace(self, FileName):
        self.m_FileName = FileName
        self.m_Replaced = True

    def WriteData(self, data):
        if self.m_Replaced == False:
            return

        strdata = self.m_FileName.encode(encoding='utf-8')

        for n in range(len(strdata)):
            data[self.m_FileNameOffset + n] = strdata[n]


# HEOReflectionProbe
class HEOReflectionProbe:

    def __init__(self, FolderNameOffset, FolderName):
        self.m_FolderNameOffset = FolderNameOffset
        self.m_FolderName = FolderName
        self.m_Hash = ""
        self.m_Replaced = False

    def MakeHash(self, RootPath):
        linkdata = bytes()

        for n in range(6):
            FileName = RootPath + "tex_reflection_cube/" + self.m_FolderName + "/reflection_cube_" + str(n) + ".png"

            with open(FileName, 'rb') as f:
                data = f.read()
                linkdata += data

        self.m_Hash = hashlib.md5(linkdata).hexdigest()

    # print("ReflectionProbe " + self.m_Hash)

    def GetHash(self):
        return self.m_Hash

    def Replace(self, FolderName):
        self.m_FolderName = FolderName
        self.m_Replaced = True

    def WriteData(self, data):
        if self.m_Replaced == False:
            return

        strdata = self.m_FolderName.encode(encoding='utf-8')

        for n in range(len(strdata)):
            data[self.m_FolderNameOffset + n] = strdata[n]


# HEOLightMap
class HEOLightMap:

    def __init__(self, FileNameOffset, FileName):
        self.m_FileNameOffset = FileNameOffset
        self.m_FileName = FileName
        self.m_Hash = ""
        self.m_Replaced = False

    def MakeHash(self, RootPath):
        FileName = RootPath + "tex/" + self.m_FileName

        with open(FileName, 'rb') as f:
            data = f.read()

        self.m_Hash = hashlib.md5(data).hexdigest()

        print(FileName + ", " + self.m_Hash)

    def GetHash(self):
        return self.m_Hash

    def Replace(self, FileName):
        self.m_FileName = FileName
        self.m_Replaced = True

    def WriteData(self, data):
        if self.m_Replaced == False:
            return

        strdata = self.m_FileName.encode(encoding='utf-8')

        for n in range(len(strdata)):
            data[self.m_FileNameOffset + n] = strdata[n]


# HEO
class HEO:

    def __init__(self, Dir, FileName):
        self.m_Dir = Dir
        self.m_FileName = FileName
        self.m_TextureList = []
        self.m_ReflectionProbeList = []
        self.m_LightMapList = []
        self.m_HEOTexCompFlag = False

    @staticmethod
    def ReadInt(data, Offset):
        return int.from_bytes(data[Offset:Offset + 4], byteorder='little')

    @staticmethod
    def ReadString(data, Offset):
        pos = data.find(0, Offset)

        if pos == -1:
            return ""

        return data[Offset: pos].decode()

    # HEOファイルの読み込み
    def LoadHEO(self, RootPath):
        FileName = RootPath + self.m_Dir + "/" + self.m_FileName

        with open(FileName, 'rb') as f:
            data = f.read()

        return data

    # HEOファイルを読み込み解析する
    def Analyse(self, RootPath):
        data = self.LoadHEO(RootPath)

        Offset = 0

        # HEAD
        head = data[Offset:3].decode()
        Offset += 4

        if head != "HEO":
            return False

        # Version
        Version = HEO.ReadInt(data, Offset)
        Offset += 4

        if Version < 26:
            return False

        # HEOTexCompを通したかのフラグ
        # バージョン32以降はテクスチャリストの後ろに移動した
        if (Version >= 21 and Version < 32):
            HEOTexCompFlag = HEO.ReadInt(data, Offset)
            Offset += 4

            self.m_HEOTexCompFlag = (HEOTexCompFlag != 0)

        # テクスチャリストバージョン
        TextureListVersion = HEO.ReadInt(data, Offset)
        Offset += 4

        if TextureListVersion != 0:
            return False

        # テクスチャリスト
        NumofTexture = HEO.ReadInt(data, Offset)
        Offset += 4

        print("NumofTexture " + str(NumofTexture))

        for n in range(NumofTexture):
            # ファイル名
            FileNameOffset = Offset

            FileName = HEO.ReadString(data, Offset)
            Offset += len(FileName) + 1

            # クランプ
            Clamp = HEO.ReadInt(data, Offset)
            Offset += 4

            # 圧縮テクスチャフラグ
            Offset += 4

            #
            self.m_TextureList.append(HEOTexture(FileNameOffset, FileName))

            print(FileName)

        # リフレクションプローブリスト
        NumofReflectionProbe = HEO.ReadInt(data, Offset)
        Offset += 4

        print("NumofReflectionProbe " + str(NumofReflectionProbe))

        for n in range(NumofReflectionProbe):
            # Skip
            Offset += 4  # UseBoxProjection
            Offset += 4 * 6  # Info
            Offset += 4  # MipmapCount
            Offset += 4  # CompFlag

            # フォルダ名
            FolderNameOffset = Offset

            FolderName = HEO.ReadString(data, Offset)
            Offset += len(FolderName) + 1

            #
            self.m_ReflectionProbeList.append(HEOReflectionProbe(FolderNameOffset, FolderName))

        # ライトマップリスト
        NumofLightMap = HEO.ReadInt(data, Offset)
        Offset += 4

        print("NumofLightMap " + str(NumofLightMap))

        for n in range(NumofLightMap):
            # ファイル名
            FileNameOffset = Offset

            FileName = HEO.ReadString(data, Offset)
            Offset += len(FileName) + 1

            # Skip
            Offset += 4  # CompFlag

            #
            self.m_LightMapList.append(HEOLightMap(FileNameOffset, FileName))

        # HEOTexCompを通したかのフラグ
        # バージョン32以降はテクスチャリストの後ろに移動した
        if (Version >= 32):
            HEOTexCompFlag = HEO.ReadInt(data, Offset)
            Offset += 4

            self.m_HEOTexCompFlag = (HEOTexCompFlag != 0)

        return True

    # テクスチャファイルのハッシュ値を生成する
    def MakeHash(self, RootPath):
        for Texture in self.m_TextureList:
            Texture.MakeHash(RootPath + self.m_Dir + "/")

        for Texture in self.m_ReflectionProbeList:
            Texture.MakeHash(RootPath + self.m_Dir + "/")

        for Texture in self.m_LightMapList:
            Texture.MakeHash(RootPath + self.m_Dir + "/")

    # マージ後のテクスチャリストの作成
    def MakeMergeTextureList(self, MergeTextureList):
        TextureIndexInHEO = 0

        for Texture in self.m_TextureList:
            Merged = False

            for mergeTexture in MergeTextureList:
                if mergeTexture.GetHash() == Texture.GetHash():
                    # 既に登録されている同一のテクスチャが見つかった
                    Texture.Replace(mergeTexture.GetFileName())

                    print("Merged InHEO(" + str(TextureIndexInHEO) + ") -> " + str(mergeTexture.GetNewIndex()))

                    Merged = True
                    break

            if Merged == False:
                NewIndex = len(MergeTextureList)
                Hash = Texture.GetHash()

                print("Set InHEO(" + str(TextureIndexInHEO) + ") -> " + str(NewIndex))

                newMergeTexture = MergeTexture(NewIndex, self, TextureIndexInHEO, Hash)
                MergeTextureList.append(newMergeTexture)

                Texture.Replace(newMergeTexture.GetFileName())

            TextureIndexInHEO += 1

    # マージ後のリフレクションプローブリストの作成
    def MakeMergeReflectionProbeList(self, MergeReflectionProbeList):
        IndexInHEO = 0

        for reflectionProbe in self.m_ReflectionProbeList:
            Merged = False

            for mergeReflectionProbe in MergeReflectionProbeList:
                if mergeReflectionProbe.GetHash() == reflectionProbe.GetHash():
                    # 既に登録されている同一のテクスチャが見つかった
                    reflectionProbe.Replace(mergeReflectionProbe.GetFolderName())

                    print("RP Merged " + str(mergeReflectionProbe.GetNewIndex()))

                    Merged = True
                    break

            if Merged == False:
                NewIndex = len(MergeReflectionProbeList)
                Hash = reflectionProbe.GetHash()

                print("RP Set " + str(NewIndex))

                newMergeReflectionProbe = MergeReflectionProbe(NewIndex, self, IndexInHEO, Hash)
                MergeReflectionProbeList.append(newMergeReflectionProbe)

                reflectionProbe.Replace(newMergeReflectionProbe.GetFolderName())

            IndexInHEO += 1

    # マージ後のライトマップリストの作成
    def MakeMergeLightMapList(self, MergeLightMapList):
        TextureIndexInHEO = 0

        for lightMap in self.m_LightMapList:
            Merged = False

            for mergeTexture in MergeLightMapList:
                if mergeTexture.GetHash() == lightMap.GetHash():
                    # 既に登録されている同一のテクスチャが見つかった
                    lightMap.Replace(mergeTexture.GetFileName())

                    print("LightMap Merged " + str(mergeTexture.GetNewIndex()))

                    Merged = True
                    break

            if Merged == False:
                NewIndex = len(MergeLightMapList)
                Hash = lightMap.GetHash()

                print("LightMap Set " + str(NewIndex))

                newMergeLightMap = MergeLightMap(NewIndex, self, TextureIndexInHEO, Hash)
                MergeLightMapList.append(newMergeLightMap)

                lightMap.Replace(newMergeLightMap.GetFileName())

            TextureIndexInHEO += 1

    # 内容を書き換えてファイルコピー
    def CopyRewriteFile(self, SrcRootPath, DestRootPath):
        data = bytearray(self.LoadHEO(SrcRootPath))

        for texture in self.m_TextureList:
            texture.WriteData(data)

        for texture in self.m_ReflectionProbeList:
            texture.WriteData(data)

        for texture in self.m_LightMapList:
            texture.WriteData(data)

        DestFileName = DestRootPath + self.m_FileName

        with open(DestFileName, 'wb') as f:
            f.write(data)

    #
    def GetFileName(self):
        return self.m_FileName

    def GetDir(self):
        return self.m_Dir

    def IsHEOTexComp(self):
        return self.m_HEOTexCompFlag

    #
    def Debug(self):
        print(self.m_Dir + ", " + self.m_FileName)


# HEOリストを作成する
def CreateHEOList(RootPath):
    # フォルダ一覧を取得する
    DirList = GetFolderList(RootPath)

    # HEOリストを作成する
    HEOList = []

    for Dir in DirList:
        # HEOファイル名一覧を取得する
        HEOFileList = glob.glob(RootPath + Dir + "/*.heo")

        for HEOFile in HEOFileList:
            heo = HEO(Dir, os.path.basename(HEOFile))
            HEOList.append(heo)

    return HEOList


# メイン
def Main():
    SrcRootPath = "./Src/"
    DestRootPath = "./Dest/"

    # HEOリストを作成する
    HEOList = CreateHEOList(SrcRootPath)

    print("NumofHEO " + str(len(HEOList)))

    for heo in HEOList:
        # ファイルの解析
        heo.Analyse(SrcRootPath)

        if heo.IsHEOTexComp() == False:
            print(heo.GetFileName() + " : imcomplete HEOTexComp")
        # return

        # テクスチャファイルのハッシュ値を計算する
        heo.MakeHash(SrcRootPath)

    # 同一テクスチャを共有化する
    MergeTextureList = []
    MergeReflectionProbeList = []
    MergeLightMapList = []

    for heo in HEOList:
        heo.MakeMergeTextureList(MergeTextureList)
        heo.MakeMergeReflectionProbeList(MergeReflectionProbeList)
        heo.MakeMergeLightMapList(MergeLightMapList)

    # テクスチャファイルのコピー
    for mergeTexture in MergeTextureList:
        mergeTexture.CopyFile(SrcRootPath, DestRootPath)

    for mergeTexture in MergeReflectionProbeList:
        mergeTexture.CopyFile(SrcRootPath, DestRootPath)

    for mergeTexture in MergeLightMapList:
        mergeTexture.CopyFile(SrcRootPath, DestRootPath)

    # ファイルを書き替えてコピー
    for heo in HEOList:
        heo.CopyRewriteFile(SrcRootPath, DestRootPath)


#
if __name__ == '__main__':
    Main()


# for launching from other scripts directly
def activate(src_root_path, dest_root_path):
    src_root_path = str(src_root_path.resolve()) + "/"
    dest_root_path = str(dest_root_path.resolve()) + "/"

    heo_list = CreateHEOList(src_root_path)
    print("NumofHEO " + str(len(heo_list)))

    for heo in heo_list:
        heo.Analyse(src_root_path)

        if not heo.IsHEOTexComp():
            print(heo.GetFileName() + " : imcomplete HEOTexComp")

        heo.MakeHash(src_root_path)

    merge_texture_list = []
    merge_reflection_probe_list = []
    merge_light_map_list = []

    for heo in heo_list:
        heo.MakeMergeTextureList(merge_texture_list)
        heo.MakeMergeReflectionProbeList(merge_reflection_probe_list)
        heo.MakeMergeLightMapList(merge_light_map_list)

    for mergeTexture in merge_texture_list:
        mergeTexture.CopyFile(src_root_path, dest_root_path)

    for mergeTexture in merge_reflection_probe_list:
        mergeTexture.CopyFile(src_root_path, dest_root_path)

    for mergeTexture in merge_light_map_list:
        mergeTexture.CopyFile(src_root_path, dest_root_path)

    for heo in heo_list:
        heo.CopyRewriteFile(src_root_path, dest_root_path)