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

# CHEPOverwriter
class CHEPOverwriter:
    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 = []

    # HEPの解析
    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))

            # 圧縮テクスチャフラグ
            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 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 Overwrite(self,FileName):
        # HEPを読む
        Analyser = CFileImageAnalyser(CFile.ReadImage(FileName))
        
        # ファイルヘッダのチェック
        head = Analyser.GetZeroTerminatedString()
        if(head != "HEP"):
            print("Error: failed to analyse head")
            return False

        # バージョン
        # HEPテクスチャ圧縮は、バージョン17から対応している
        Version = Analyser.GetInt()
        if(Version < 17):
            print("Error: HEP ver.%d doesn't support texture compression. Use ver.17 or later." %(Version))
            return False

        # HEPTexCompを通したかのフラグ
        HEPTexCompFlagOffset = Analyser.m_DataOffset
        HEPTexCompFlag = Analyser.GetInt()    

        # フラグを書き換える
        Analyser.SetInt(1,HEPTexCompFlagOffset)
        
        # テクスチャリスト  
        if(not self.AnalyseTextureList(Analyser,Version)):
            print("Error: failed to analyse TextureList")
            return False     

        # 圧縮テクスチャファイルの存在をチェックし指定のヘッダ位置のフラグを書き替える
        if(not self.AnalyseTextureCompression(Version,FileName,Analyser)): # 通常のテクスチャの圧縮フラグを確認する
            print("Error: failed to analyse TextureCompression")
            return False

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

        return True

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

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

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

# 関数実行 
Main()