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):
		# HEOバージョン35からフォルダ表記が 『000』ではなく『0』になっているので読み替える
		FolderName = ""

		if(self.m_HEO.GetVersion() >= 35):
			FolderName = str(self.m_NewIndex)
		else:
			FolderName = format(str(self.m_NewIndex), '0>3')

		return FolderName

	def GetSrcFolderName(self):
		# HEOバージョン35からフォルダ表記が 『000』ではなく『0』になっているので読み替える
		FolderName = ""

		if(self.m_HEO.GetVersion() >= 35):
			FolderName = str(self.m_IndexInHEO)
		else:
			FolderName = format(str(self.m_IndexInHEO), '0>3')

		return FolderName

	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):
		# ライトマップはRGBMかつetc2の時だけ圧縮可能になっているのでその分も書き出す
		FolderNameList = [ "tex", "tex_etc" ]
		ExtList = [ "png", "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)


# 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_Version = 0
		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

		print("FileName: %s, HEO Version: %d" %(self.m_FileName, Version))

		if Version < 26:
			return False
		
		self.m_Version = Version

		# 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 GetVersion(self):
		return self.m_Version	

	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)

#
Main()