#!/usr/bin/env python3

from argparse import ArgumentParser, FileType
from io import TextIOWrapper
from os import path, remove, system
from struct import pack
from sys import argv, platform


def img2tdx(tds, tdx, header=None, imgPath="."):
	"""
	Converts images to a TDX file, inspired by 3DS T3X files

	Parameters
	----------
	tds
		Can be a file or a list of file paths, the first item can optionally be arguments for grit
	output
		The file to output to
	header
		(optional) Can be a file to output a C header to
	imgPath
		(optional) Path to where the images are
	"""

	if platform == "win32":
		if system("where /q grit") != 0:
			raise Exception("GRIT must be in your PATH")
	else:
		if system("which grit > /dev/null") != 0:
			raise Exception("GRIT must be in your PATH")

	gritArgs = "-ftr -fh!"

	if type(tds) == TextIOWrapper:
		paths = tds.read().split("\n")
		tdsPath = path.dirname(tds.name)
	else:
		paths = tds
		tdsPath = imgPath

	if not path.exists(paths[0]):
		gritArgs += " " + paths[0]
		paths = paths[1:]

	# Start with a hash to comment a file
	paths = [f"{tdsPath}/{item}" for item in paths if not (item == "" or item[0] == "#")]

	files = []

	# Read files and convert to GRF
	for filePath in paths:
		if filePath.endswith(".grf"):
			with open(filePath, "rb") as grf:
				files.append(grf.read())
		else:
			grfPath = filePath[:filePath.rfind(".")] + ".grf"
			if system(f'grit "{filePath.replace("/", path.sep)}" {gritArgs} -o"{grfPath}"') == 0:
				with open(grfPath, "rb") as grf:
					files.append(grf.read())
				remove(grfPath)
			else:
				raise Exception("Error while running GRIT")

	# Write tdx
	tdx.write(b".TDX")  # Format magic
	tdx.write(pack("<L", len(files)))  # File count
	tdx.write(b"\x00" * 8)  # Padding

	# Position of data, starts after the file list
	position = 0x10 + len(files) * 8

	# Write file list
	for file in files:
		# Position of data, Size of data, Image width, Image height
		tdx.write(pack("<LL", position, len(file)))
		position += len(file)

	# Write data
	for file in files:
		tdx.write(file)

	if header:
		header.write("/* Generated by %s */\n#pragma once\n\n" % argv[0])
		for i, file in enumerate(paths):
			header.write("#define sprites_%s_idx %d\n" % (path.splitext(path.basename(file))[0], i))


if __name__ == "__main__":
	parser = ArgumentParser(description="Converts images with GRIT and packs them together")
	parser.add_argument("tds", metavar="input.tds", type=FileType("r"), help="tds to include from")
	parser.add_argument("-o", "--output", metavar="output.tdx", default="output.tdx", type=FileType("wb"), help="file to output to")
	parser.add_argument("--header", metavar="output.h", type=FileType("w"), help="C header of image indexes")

	args = parser.parse_args()

	img2tdx(args.tds, args.output, args.header)
