#!/usr/bin/env python3

import argparse
import concurrent.futures
import json
import os
import pathlib
import re
import shutil
import subprocess

ROOT = os.path.dirname(os.path.abspath(__file__))
imagecompare = f'{ROOT}/imagecompare'
WORKDIR = f'{ROOT}/work'
REPORT_ROOT = f'{ROOT}' # FIXME should be workdir?

# TODO rename to <map>_<description>_<issue(opt)>
saves = []
for (_, _, files) in os.walk(os.path.join(ROOT, 'save')):
	for file in files:
		saves.append(file.removesuffix('.sav').removeprefix('rendertest_'))

channels = {
	'full': '',
	'basecolor': 'basecolor',
	'emissive': 'emissive',
	'nshade': 'nshade',
	'ngeom': 'ngeom',
	'lighting': 'lighting',
	'direct': 'direct',
	'indirect': 'indirect',
	'indirect_specular': 'indirect_spec',
	'indirect_diffuse': 'indirect_diff',
}

def test_list(arg: str) -> [str]:
	items = arg.split(',')
	tests = []
	for item in items:
		r = re.compile(item)
		for save in saves:
			if r.search(save):
				tests.append(save)
	if not tests:
		raise argparse.ArgumentTypeError(f'No tests match {item}. Available tests are: {saves}')
	return tests

parser = argparse.ArgumentParser(description='Generate scripts and makefiles for rendertest')
#parser.add_argument('--script', '-s', type=argparse.FileType('w'), help='Console script for generating images')
parser.add_argument('--tests', '-t', type=test_list, default=saves, help='Run only these tests')

# TODO how to check that the dir is valid? presence of xash3d executable and valve dir?
parser.add_argument('--xash-dir', '-x', type=str, default=os.getcwd(), help='Path to xash3d-fwgs installation directory')

# TODO parse commands in type=.. function
parser.add_argument('command', type=str, default=None, help='Action to perform')
args = parser.parse_args()

def make_script(file, tests: [str]):
	header = '''sv_cheats 1
developer 0
m_ignore 1
cl_showfps 0
scr_conspeed 1000000
con_notifytime 0
hud_draw 0
r_speeds 0
rt_debug_fixed_random_seed 31337

'''

	print(f'Generating script {file.name}')
	file.write(header)

	for test in tests:
		screenshot_base = 'rendertest/'
		file.write(f'load rendertest_{test}\n')
		file.write(f'wait 4; echo DONE WAIT4; playersonly; wait 11\n')
		# for i in range(13):
		# 	file.write(f'echo FRAME {i+4}; wait 1;\n')
		for channel, display in channels.items():
			file.write(f'rt_debug_display_only "{display}"; screenshot {screenshot_base}{test}_{channel}.tga; wait 1\n')
		file.write('\n')

	file.write('quit\n')

# if args.script:
# 	print(f'Generating script {args.script.name}')
# 	make_script(args.script, args.tests)

def mkdir_p(path: str):
	pathlib.Path(path).mkdir(parents=True, exist_ok=True)

def compile():
	subprocess.run(['make', 'imagecompare'], cwd=ROOT, check=True)

def copy_assets():
	print('Copying assets')
	shutil.copy2(src=f'{ROOT}/vulkan_debug.wad', dst=f'{args.xash_dir}/valve/')
	shutil.copytree(src=f'{ROOT}/maps', dst=f'{args.xash_dir}/valve/maps/', dirs_exist_ok=True)
	shutil.copytree(src=f'{ROOT}/save', dst=f'{args.xash_dir}/valve/save/', dirs_exist_ok=True)
	with open(f'{args.xash_dir}/rendertest.script', 'w') as script:
		make_script(script, args.tests)

def render():
	print('Running xash3d')
	mkdir_p(f'{args.xash_dir}/valve/rendertest')
	env = os.environ.copy()
	env['RADV_PERFTEST'] = 'rt'
	env['LD_LIBRARY_PATH'] = '.'
	subprocess.run([f'{args.xash_dir}/xash3d', '-ref', 'vk',
		'-nowriteconfig', '-nosound', '-log',
		'-width', '1280', '-height', '800',
		'+exec', 'rendertest.script'],
		env=env, check=True)

def compare_one(test: str, channel: str, image_base: str, image_gold: str, image_test: str, image_diff: str, image_flip: str):
	result = subprocess.run([imagecompare, image_gold, image_test, image_diff], text=True, capture_output=True)
	match result.returncode:
		case 0:
			pass
		case 1:
			raise Exception(f'FATAL imagecompare error: TBD')
		case 2:
			print(f'ERROR: {image_base} differ by more than threshold')
		case _:
			raise Exception(f'Unexpected imagecompare return code {result.returncode}')
	ret = json.loads(result.stdout)
	ret['test'] = test
	ret['channel'] = channel
	ret['image_gold'] = os.path.relpath(image_gold, REPORT_ROOT)
	ret['image_test'] = os.path.relpath(image_test, REPORT_ROOT)
	ret['image_diff'] = os.path.relpath(image_diff, REPORT_ROOT)
	ret['image_flip'] = os.path.relpath(image_flip, REPORT_ROOT)
	return ret

def compare():
	mkdir_p(WORKDIR)
	screenshot_base = f'{args.xash_dir}/valve/rendertest'
	diffs = []
	with concurrent.futures.ThreadPoolExecutor() as executor:
		for test in args.tests:
			for channel, _ in channels.items():
				image_base = f'{test}_{channel}'
				image_test = f'{screenshot_base}/{image_base}.tga'
				image_gold = f'{ROOT}/gold/{image_base}.png'
				image_diff = f'{WORKDIR}/{image_base}_diff.png'
				image_flip = f'{WORKDIR}/{image_base}_flip.gif'

				diffs.append(executor.submit(compare_one, test, channel, image_base, image_gold, image_test, image_diff, image_flip))

				executor.submit(subprocess.run, ['convert',
					'(', image_gold, '-bordercolor', 'gold', '-border', '2x2', '-gravity', 'SouthWest', '-font', 'Impact', '-pointsize', '24', '-fill', 'gold', '-stroke', 'black', '-annotate', '0', 'GOLD', ')',
					'(', image_test, '-bordercolor', 'white', '-border', '2x2', '-fill', 'white', '-annotate', '0', 'TEST', ')',
					'-loop', '0', '-set', 'delay', '100', image_flip], check=True)

	results = [diff.result() for diff in diffs]
	# json.dump(results, open(f'{WORKDIR}/data.json', 'w'))
	jsons = json.dumps(results)
	with open(f'{WORKDIR}/data.js', 'w') as js:
		js.write(f'"use strict";\nconst data = {jsons};\n')

def command_compare():
	compile()
	compare()

def command_png():
	screenshot_base = f'{args.xash_dir}/valve/rendertest'
	new_gold_base = f'{WORKDIR}/gold'
	mkdir_p(new_gold_base)
	with concurrent.futures.ThreadPoolExecutor() as executor:
		for test in args.tests:
			for channel, _ in channels.items():
				image_base = f'{test}_{channel}'
				image_test = f'{screenshot_base}/{image_base}.tga'
				image_new_gold = f'{new_gold_base}/{image_base}.png'

				print(f'{image_new_gold}')
				executor.submit(subprocess.run, ['convert', image_test, image_new_gold], check=True)

def command_run():
	compile()
	copy_assets()
	render()
	compare()

def command_render():
	compile()
	copy_assets()
	render()

# TODO dict
match args.command:
	case 'run':
		command_run()
	case 'compare':
		command_compare()
	case 'render':
		command_render()
	case 'png':
		command_png()
	case _:
		raise Exception(f'Unknown command {args.command}')

# TODO:
# - settings object to pass as an argument
# - HTML report