import sys
import struct
import glob

# CFile
class CFile:
    def ReadImage(FileName):
        with open(FileName,'rb') as f:
            Data=f.read()

        return Data

    def WriteImage(FileName,Data):
        with open(FileName, 'wb') as f:
            f.write(Data)

    def FindFile(FilePath):
        FindFileNameList = glob.glob(FilePath)
        return (len(FindFileNameList) > 0)

# LFileName
class LFileName:
    def DeleteFileName(FileName):
        IndexOfYen = FileName.rfind('\\')
        IndexOfSlash = FileName.rfind('/')
        
        if(IndexOfYen != -1):
            DirName = FileName[:IndexOfYen + 1]
            return DirName
        elif(IndexOfSlash != -1):
            DirName = FileName[:IndexOfSlash + 1]
            return DirName
        else:
            DirName = ""
            return DirName

    def AddPunct(FileName):
        IndexOfYen = FileName.rfind('\\')
        IndexOfSlash = FileName.rfind('/')
        
        if(IndexOfYen != -1):
            Punct = '\\'
            return Punct
        elif(IndexOfSlash != -1):
            Punct = '/'
            return Punct
        else:
            Punct = ""
            return Punct

    def ChangeExtension(FileName,Extension):
        IndexOfExtensionDot = FileName.rfind('.')
        RenamedFileName = ""

        if(IndexOfExtensionDot != -1):
            RenamedFileName = FileName[:IndexOfExtensionDot + 1] + Extension
        else:
            RenamedFileName = FileName + "." + Extension

        return RenamedFileName

# CFileImageAnalyser
class CFileImageAnalyser:
    def __init__(self,Data):
        self.m_DataOffset=0
        self.m_Data=bytearray(Data)
        self.m_DataEnd=len(self.m_Data)

    def GetZeroTerminatedString(self):
        # 開始オフセット
        StartOffset = self.m_DataOffset
        
        # 0が来るまでループする
        while(True):
            if(self.m_Data[self.m_DataOffset] == 0):
                break
            else :
                self.m_DataOffset+=1

        # 終了オフセット
        EndOffset =  self.m_DataOffset     

        val = self.m_Data[StartOffset:EndOffset].decode()
        self.m_DataOffset+=1
        return val

    def GetInt(self):
        val = int.from_bytes(self.m_Data[self.m_DataOffset:self.m_DataOffset + 4], byteorder='little')
        self.m_DataOffset += 4
        return val

    def SetInt(self,SetVal,StartOffset):
        ByteVal = bytearray((SetVal).to_bytes(4, byteorder='little'))
        
        self.m_Data[StartOffset + 0] = ByteVal[0]
        self.m_Data[StartOffset + 1] = ByteVal[1]
        self.m_Data[StartOffset + 2] = ByteVal[2]
        self.m_Data[StartOffset + 3] = ByteVal[3]

    def GetFloat(self):
        val = struct.unpack('<f',self.m_Data[self.m_DataOffset:self.m_DataOffset + 4])
        self.m_DataOffset += 4
        return val

    def SetByte(self,SetVal,StartOffset):
        ByteVal = bytearray((SetVal).to_bytes(1, byteorder='little'))
        self.m_Data[StartOffset + 0] = ByteVal[0]

    def GetByte(self):
        val = self.m_Data[self.m_DataOffset]
        self.m_DataOffset += 1

        return val

# ETextureCompressionType
class ETextureCompressionType():
    PVRTC = 0
    ETC2 = 1
    ASTC = 2
    DXT = 3

# CHEOOverwriter
class CHEOOverwriter:
    def __init__(self):
        self.m_TextureCompressionInfoPointerList = []
        self.m_TextureFileNameList = []

        self.m_CubemapTextureCompressionInfoPointerList = []
        self.m_ReflectionProbeFolderNameList = []
        self.m_ReflectionProbeNum = 0

        self.m_LightmapFileNameList = []
        self.m_LightmapCompressionInfoPointerList = []

    # HEOの解析
    def AnalyseTextureList(self,Analyser,Version):

        # 個数
        NumOf = Analyser.GetInt()

        # テクスチャリスト
        for n in range(NumOf):
            # ファイル名
            FileName = Analyser.GetZeroTerminatedString()
            self.m_TextureFileNameList.append(FileName)
            #print("FileName: %s" %(FileName))

            # クランプ
            Clamp = False

            if (Version >= 1):
                ClampType = 0
                ClampType = Analyser.GetInt()

                Clamp = (ClampType == 1)

            # 圧縮テクスチャフラグ
            TextureCompressionSupportedList = [0,0,0,0]
            TextureCompressionSupportedPointerList = [0,0,0,0]

            if (Version >= 5):
                for m in range(4):
                    # ここに圧縮テクスチャフラグがBYTE型で存在する
                    TextureCompressionSupportedPointerList[m] = Analyser.m_DataOffset

                    TextureCompressionSupportedList[m] = Analyser.GetByte()

            # テクスチャ圧縮フラグリストに登録
            self.m_TextureCompressionInfoPointerList.append(TextureCompressionSupportedPointerList)

        return True
    
    def AnalyseReflectionProbe(self,Analyser,Version):
        # 使用するリフレクションプローブの総数を読み取る
        ReflectionProbeNum = 0
        if (Version >= 11):
            ReflectionProbeNum = Analyser.GetInt()

        # リフレクションプローブリスト
        if (Version >= 12):
            for i in range(ReflectionProbeNum):
                # Box Projection を使用するかどうか
                UseBoxProjection = Analyser.GetInt()

                # リフレクションプローブの中心位置(オフセット)
                ReflectionProbeCenter_X = Analyser.GetFloat()
                ReflectionProbeCenter_Y = Analyser.GetFloat()
                ReflectionProbeCenter_Z = Analyser.GetFloat()

                # リフレクションプローブのサイズ
                ReflectionProbeSize_X = Analyser.GetFloat()
                ReflectionProbeSize_Y = Analyser.GetFloat()
                ReflectionProbeSize_Z = Analyser.GetFloat()

                if (Version >= 16):
                    # リフレクションプローブのMipmapCountを読み込む
                    MipMapCount = Analyser.GetInt()

                    # 圧縮テクスチャフラグ
                    TextureCompressionSupportedList = [0,0,0,0]
                    TextureCompressionSupportedPointerList = [0,0,0,0]

                    if (Version >= 5):
                        for m in range(4):
                            # ここに圧縮テクスチャフラグがBYTE型で存在する
                            TextureCompressionSupportedPointerList[m] = Analyser.m_DataOffset

                            TextureCompressionSupportedList[m] = Analyser.GetByte()

                    # テクスチャ圧縮フラグリストに登録
                    self.m_CubemapTextureCompressionInfoPointerList.append(TextureCompressionSupportedPointerList)

                    # フォルダ名
                    if (Version >= 26):
                        FolderName = Analyser.GetZeroTerminatedString()
                        #print("FolderName: %s" %(FolderName))
                        self.m_ReflectionProbeFolderNameList.append(FolderName)
                    else:
                        self.m_ReflectionProbeFolderNameList.append(str(i))

        #
        self.m_ReflectionProbeNum = ReflectionProbeNum

        return True
    
    def AnalyseLightmapList(self,Analyser,Version):
        Numof = Analyser.GetInt()

        #
        if (Version >= 26):
            for n in range(Numof):
                # ファイル名
                FileName = Analyser.GetZeroTerminatedString()
                #print("FileName: %s" %(FileName))
                self.m_LightmapFileNameList.append(FileName)

                # 圧縮テクスチャフラグ
                TextureCompressionSupportedList = [0,0,0,0]
                TextureCompressionSupportedPointerList = [0,0,0,0]

                if (Version >= 5):
                    for m in range(4):
                        # ここに圧縮テクスチャフラグがBYTE型で存在する
                        TextureCompressionSupportedPointerList[m] = Analyser.m_DataOffset

                        TextureCompressionSupportedList[m] = Analyser.GetByte()

                    # テクスチャ圧縮フラグリストに登録
                    self.m_LightmapCompressionInfoPointerList.append(TextureCompressionSupportedPointerList)

        # シーンにリアルタイムのDirectionalLightを使用しているか
        # 使用している場合、LightMapによるライティングとPBRによるライティングが合成される
        UseDirectionlLight = False
        UseDirectionlLightFlag = 0
        if (Version >= 18):
            UseDirectionlLightFlag = Analyser.GetInt()
            if(UseDirectionlLightFlag == 1):
                UseDirectionlLight = True

        # LightMapのエンコードタイプ
        if (Version >= 22):
            LightMapEncodeType = Analyser.GetInt()

        return True
    
    # 圧縮テクスチャの存在を調べて圧縮フラグを書き換える
    def AnalyseTextureCompression(self,Version,FileName,Analyser):
        #
        NumofTexComp = [0,0,0,0]
        DirName = LFileName.DeleteFileName(FileName)
        Punct = LFileName.AddPunct(FileName)
        
        # DirNameとPunctが空の場合はカレントディレクトリ(./)を示している
        if(DirName == "" or Punct == ""):
            DirName = "./"
            Punct = "/"

        #
        for TextureIndex in range(len(self.m_TextureCompressionInfoPointerList)):
            TextureCompressionInfoPointer = self.m_TextureCompressionInfoPointerList[TextureIndex]

            #
            for n in range(4):
                SubFolderName = ""
                Extension = "pvr"

                if(n == ETextureCompressionType.PVRTC):
                    SubFolderName = "tex_pvrtc"
                elif(n == ETextureCompressionType.ETC2):
                    SubFolderName = "tex_etc"
                elif(n == ETextureCompressionType.ASTC):
                    SubFolderName = "tex_astc"
                elif(n == ETextureCompressionType.DXT):
                    SubFolderName = "tex_dxt"
                    Extension = "dds"
                else:
                    return False
                
                # 圧縮テクスチャのパスを生成
                TexFileName = DirName + SubFolderName + Punct + LFileName.ChangeExtension(self.m_TextureFileNameList[TextureIndex] , Extension)
                #print("TexFileName: %s" %(TexFileName))
                
                # 圧縮フラグをファイルの検索結果に応じて書き換える
                if(not CFile.FindFile(TexFileName)): 
                    Analyser.SetByte(0,TextureCompressionInfoPointer[n]) # 圧縮テクスチャが存在しない
                else:
                    Analyser.SetByte(1,TextureCompressionInfoPointer[n]) # 圧縮テクスチャが存在する
                    NumofTexComp[n]+=1

        print("ExistTextureCompression %d %d %d %d" %(NumofTexComp[0],NumofTexComp[1],NumofTexComp[2],NumofTexComp[3]))
                
        return True

    def AnalyseReflectionProbeCompression(self,Version,FileName,Analyser):
        #
        NumofTexComp = [0,0,0,0]
        DirName = LFileName.DeleteFileName(FileName)
        Punct = LFileName.AddPunct(FileName)
        
        # DirNameとPunctが空の場合はカレントディレクトリ(./)を示している
        if(DirName == "" or Punct == ""):
            DirName = "./"
            Punct = "/"

        #
        for ReflectionProbeIndex in range(self.m_ReflectionProbeNum):
            TextureCompressionInfoPointer = self.m_CubemapTextureCompressionInfoPointerList[ReflectionProbeIndex]
            IsCompressList = [True,True,True,True]

            # リフレクションプローブはテクスチャ6枚で1セット
            for TextureIndex in range(6):
                #
                for n in range(4):
                    SubFolderName = ""
                    Extension = "pvr"

                    if(n == ETextureCompressionType.PVRTC):
                        SubFolderName = "tex_reflection_cube_pvrtc"
                    elif(n == ETextureCompressionType.ETC2):
                        SubFolderName = "tex_reflection_cube_etc"
                    elif(n == ETextureCompressionType.ASTC):
                        SubFolderName = "tex_reflection_cube_astc"
                    elif(n == ETextureCompressionType.DXT):
                        SubFolderName = "tex_reflection_cube_dxt"
                        Extension = "dds"
                    else:
                        return False
                
                    # 圧縮テクスチャのパスを生成
                    TexFileName = DirName + SubFolderName + Punct + self.m_ReflectionProbeFolderNameList[ReflectionProbeIndex] + Punct
                    TexFileName += LFileName.ChangeExtension( "reflection_cube_" + str(TextureIndex), Extension)
                    
                    # 圧縮テクスチャが一つでも存在しないのであれば、そのリフレクションプローブは圧縮テクスチャを取り扱わないものとする
                    if(not CFile.FindFile(TexFileName)):
                        IsCompressList[n] =False

                    #print("TexFileName: %s" %(TexFileName))
                
            # 圧縮フラグをファイルの検索結果に応じて書き換える 
            for n in range(4):
                IsCompress = IsCompressList[n]

                if(IsCompress): # 圧縮テクスチャが存在する
                    NumofTexComp[n] += 1
                    Analyser.SetByte(1,TextureCompressionInfoPointer[n])
                else: # 圧縮テクスチャが存在しない
                    Analyser.SetByte(0,TextureCompressionInfoPointer[n])
      
        print("ExistCubemapCompression %d %d %d %d" %(NumofTexComp[0],NumofTexComp[1],NumofTexComp[2],NumofTexComp[3]))
        return True

    def AnalyseLightmapCompression(self,Version,FileName,Analyser):
        #
        NumofTexComp = [0,0,0,0]
        DirName = LFileName.DeleteFileName(FileName)
        Punct = LFileName.AddPunct(FileName)
        
        # DirNameとPunctが空の場合はカレントディレクトリ(./)を示している
        if(DirName == "" or Punct == ""):
            DirName = "./"
            Punct = "/"

        #
        for LightmapIndex in range(len(self.m_LightmapCompressionInfoPointerList)):
            LightmapCompressionInfoPointer = self.m_LightmapCompressionInfoPointerList[LightmapIndex]

            #
            for n in range(4):
                SubFolderName = ""
                Extension = "pvr"

                if(n == ETextureCompressionType.PVRTC):
                    SubFolderName = "tex_pvrtc"
                elif(n == ETextureCompressionType.ETC2):
                    SubFolderName = "tex_etc"
                elif(n == ETextureCompressionType.ASTC):
                    SubFolderName = "tex_astc"
                elif(n == ETextureCompressionType.DXT):
                    SubFolderName = "tex_dxt"
                    Extension = "dds"
                else:
                    return False
                
                # 圧縮テクスチャのパスを生成
                TexFileName = DirName + SubFolderName + Punct + LFileName.ChangeExtension(self.m_LightmapFileNameList[LightmapIndex] , Extension)
                #print("TexFileName: %s" %(TexFileName))
                
                # 圧縮フラグをファイルの検索結果に応じて書き換える
                if(not CFile.FindFile(TexFileName)): 
                    Analyser.SetByte(0,LightmapCompressionInfoPointer[n]) # 圧縮テクスチャが存在しない
                else:
                    Analyser.SetByte(1,LightmapCompressionInfoPointer[n]) # 圧縮テクスチャが存在する
                    NumofTexComp[n]+=1

        print("ExistLightmapCompression %d %d %d %d" %(NumofTexComp[0],NumofTexComp[1],NumofTexComp[2],NumofTexComp[3]))
                
        return True
    
    #
    def Overwrite(self,FileName):
        # HEOを読む
        Analyser = CFileImageAnalyser(CFile.ReadImage(FileName))
        
        # ファイルヘッダのチェック
        head = Analyser.GetZeroTerminatedString()
        if(head != "HEO"):
            print("Error: failed to analyse head")
            return False

        # バージョン
        Version = Analyser.GetInt()
        if(Version < 7):
            print("Error: failed to analyse HEO version")
            return False

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

            # フラグを書き換える
            Analyser.SetInt(1,HEOTexCompFlagOffset)
        
        # テクスチャリスト  
        # Python版のHEOTexCompではVer.26以降しか想定していない
        # Ver.26以前のものは、これまで通りHEOTexComp.exeを使用する  
        if (Version >= 26):
            TextureListVersion = Analyser.GetInt()
            if(TextureListVersion != 0):
                print("Error: failed to analyse TextureListVersion")
                return False

            # 通常の2Dテクスチャ
            if(not self.AnalyseTextureList(Analyser,Version)):
                print("Error: failed to analyse TextureList")
                return False

            # リフレクションプローブ
            if(not self.AnalyseReflectionProbe(Analyser,Version)):
                print("Error: failed to analyse ReflectionProbe")
                return False

            # ライトマップ
            if(not self.AnalyseLightmapList(Analyser,Version)):
                print("Error: failed to analyse LightmapList")
                return False

            # HEOTexCompを通したかのフラグ
            # バージョン32以降はテクスチャリストの後ろに移動した
            if(Version >= 32):
                HEOTexCompFlagOffset = Analyser.m_DataOffset
                HEOTexCompFlag = Analyser.GetInt()

                # フラグを書き換える
                Analyser.SetInt(1,HEOTexCompFlagOffset)


        # 圧縮テクスチャファイルの存在をチェックし指定のヘッダ位置のフラグを書き替える
        if(not self.AnalyseTextureCompression(Version,FileName,Analyser)): # 通常のテクスチャの圧縮フラグを確認する
            print("Error: failed to analyse TextureCompression")
            return False
        if(not self.AnalyseReflectionProbeCompression(Version,FileName,Analyser)): # リフレクションプローブキューブマップの圧縮フラグを確認する
            print("Error: failed to analyse ReflectionProbeCompression")
            return False
        if(not self.AnalyseLightmapCompression(Version,FileName,Analyser)): # ライトマップの圧縮フラグを確認する
            print("Error: failed to analyse LightmapCompression")
            return False

        # 圧縮フラグを更新したHEOを書き出す
        CFile.WriteImage(FileName,Analyser.m_Data)

        return True

# メイン
def Main():
    # コマンドライン引数
    args=sys.argv
    if(len(args) <2):
        print("Error : Please specify the path to HEO in the command line argument.")
        return
    
    FileName = args[1]

    # 圧縮テクスチャの存在を調べて圧縮フラグを書き換える
    HEOOverwriter=CHEOOverwriter()
    if  not HEOOverwriter.Overwrite(FileName):
        print("HEOTexComp failed to execute with error. : %s\n" % (FileName))
        return


    # フラグの書き換えに成功
    print("Succeeded\n")

# 関数実行 
Main()