# parse text maps
import helpers.utils as utils
from parsers.text import Textmap

import re
import os
import glob
import shutil
import json
import subprocess
from rich import print
from io import BytesIO
from rich.markdown import Markdown
from kaitaistruct import KaitaiStream
from rich.progress import Progress, SpinnerColumn, TimeElapsedColumn, Console

# TODO: use config
map_files = {
	"CHS": 26692920,
	"CHT": 27251172,
	"DE": 25181351,
	"EN": 25776943,
	"ES": 20618174,
	"FR": 25555476,
	"ID": 30460104,
	"JP": 32244380,
	"KR": 22299426,
	"PT": 23331191,
	"RU": 21030516,
	"TH": 32056053,
	"VI": 34382464,
	"IT": 27270675,
	"TR": 21419401
}

ASCII = re.compile(r"[^\u0020-\u007E]")

def extract_text_maps(requested):
	game_target = utils.get_config("gameTarget")

	if not game_target:
		print("[red3]No game target !\nDefine it with --set_target first[/]")
		return

	console = Console(record=True)

	for lang in requested:
		blk_file = utils.get_path(f"blk/{map_files[lang]}.blk")

		# find xor key
		# actually xor key is annoying to find so manually defined for now, need to update at every version
		xor_key = 99
		print(f"XOR key : {xor_key}")

		# extract mhy data
		with console.status("[cyan]Extracting text data...", spinner="dots") as status:
			process = subprocess.Popen([utils.get_config("studioPath"), blk_file, utils.get_path("bin/"), "--game", "GI", "--types", "MiHoYoBinData", "--key", str(int(str(xor_key), 16))], stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
			process.wait()


		# xor each file
		# doesn't work, encoding is missing and headers too
		files = glob.glob(utils.get_path("bin/MiHoYoBinData") + "/*")
		print(f"[cyan] [b]{len(files)}[/] assets extracted !")
		# pos = 0

		# for file in files:
		# 	xor = int(str(xor_key), 16)
		# 	pos += 1
		# 	print(f"[{pos}/{len(files)}] Xoring {os.path.basename(file)}", end="\r")
		# 	with open(file, "r+b") as f:
		# 		data = f.read()

		# 		if len(data) == 0:
		# 			f.close()
		# 			continue

		# 		data = ASCII.sub("", "".join([chr(b ^ xor) for b in data]))
				
		# 		f.seek(0)
		# 		f.write(data.encode("utf-8"))
		# 		f.truncate()

		# 		f.close()

		# parse each file
		output = {}
		pos = 0

		with Progress(
			SpinnerColumn(),
			*Progress.get_default_columns(),
			TimeElapsedColumn(),
			console=console,
			transient=False
			) as progress:

			task = progress.add_task(f"[cyan]Building ({lang}) text map...", total=len(files))

			for file in files:
				pos += 1
				progress.update(task, advance=1)
				# print(f"[{pos}/{len(files)}] Parsing {os.path.basename(file)}", end="\r")

				# could store xor result so we don't reopen file ?
				with open(file, "rb") as f:
					data = f.read()

					if len(data) == 0:
						f.close()
						continue

					try:
						stream = KaitaiStream(BytesIO(data))
						obj = Textmap(stream)
					except UnicodeDecodeError:
						print("Failed to parse, is xor ok ?")
						break

					for block in obj.textmap:
						if block.string.data != "":
							output[str(block.hash.value)] = block.string.data
					
					f.close()

		out_path = utils.get_path(f"output/TextMap/TextMap{lang}.json")
		os.makedirs(os.path.dirname(out_path), exist_ok=True)

		# use ai file to sort like dim instead :/
		output = dict(sorted(output.items(), key=lambda x: int(x[0])))

		with open(out_path, "w", encoding='utf-8') as f:
			json.dump(output, f, indent=4, ensure_ascii=False)
			f.close()

		filesize = os.path.getsize(out_path)

		shutil.rmtree(utils.get_path("bin/MiHoYoBinData"))

		# print(f"Extracted {lang} map !")

		progress.log(Markdown("## Done !"))
		progress.log(Markdown(f"# Entries : {len(output)} | Size : `{round(filesize / 1048576, 2)} Mb`"))

# extract_text_maps(["FR"])