from pathlib import Path
from shutil import rmtree
from sys import argv
from argparse import ArgumentParser, Namespace
import tomllib as toml
import re

MX_DISCARD_PATHS = [
	"cmake",
	".mxproject",
	"CMakePresets.json",
]

SOURCE_EXTS = [
	"*.c",
	"*.s"
]
HEADER_EXTS = [
	"*.h"
]

CMAKE_EXCLUDE_SOURCES = [
	"*_template.c"
]
CMAKE_TARGET = ""
ASSUME_YES = False
INJECT_HANDLERS = []
NO_MY_MAIN = False


def mx_cleanup(path: Path):
	deletions: list[Path] = []

	for disc in MX_DISCARD_PATHS:
		disc_path = path / disc
		if disc_path.exists():
			deletions.append(disc_path)

	if len(deletions) == 0:
		return

	print("%d files selected for deletion:"%(len(deletions)))
	for disc in deletions:
		print(disc)

	if not ASSUME_YES:
		print("Accept? (y/N)")
		sel = input()
		if sel.lower() != "y":
			print("Aborting...")
			exit(1)
	
	for disc in deletions:
		disc.unlink() if disc.is_file() else rmtree(disc)

def recurse_cmakelists(path: Path):
	# Detect relevant files
	dirs: list[Path] = []
	sources: list[Path] = []
	has_headers: bool = False

	for child in sorted(path.iterdir()):
		if child.is_dir():
			dirs.append(child)
		elif any(child.match(ext) for ext in SOURCE_EXTS) and \
			not any(child.match(exclude) for exclude in CMAKE_EXCLUDE_SOURCES):
			sources.append(child)
		elif any(child.match(ext) for ext in HEADER_EXTS):
			has_headers = True

	# Generate CMakeLists content
	lines: list[str] = []

	if has_headers:
		lines.append("target_include_directories(%s PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})"%(CMAKE_TARGET))
	if len(sources) != 0:
		lines.append("target_sources(%s PUBLIC"%(CMAKE_TARGET))
		lines += ["\t%s"%(src.name) for src in sources]
		lines.append(")")
	lines += ["add_subdirectory(%s)"%(directory.name) for directory in dirs]

	lists_path = path / "CMakeLists.txt"
	with open(lists_path, "w") as of:
		of.write("\n".join(lines)+"\n")

	# Recurse
	for directory in dirs:
		recurse_cmakelists(directory)

def root_cmakelists(path: Path):
	lines: list[str] = []
	lists_path = path / "CMakeLists.txt"

	# Define custom libraries
	lines.append("add_library(%s STATIC)"%(CMAKE_TARGET))
	lines.append("target_compile_definitions(%s PRIVATE ${COMPILER_DEFINES})"%(CMAKE_TARGET))
	lines.append("target_compile_options(${EXECUTABLE} PRIVATE -Os -ffunction-sections -fdata-sections -g -flto)")
	lines.append("target_link_options(${EXECUTABLE} PRIVATE -T ${LINKER_SCRIPT} --specs=nano.specs -Wl,--gc-sections -flto)")

	# Read previously generated contents
	with open(lists_path, "r") as _if:
		lines += [line.rstrip("\n") for line in _if]
	
	# Save final result
	with open(lists_path, "w") as of:
		of.write("\n".join(lines)+"\n")

def inject_handlers_file(path: Path):
	with open(path, "r") as _if:
		content = _if.read()

	patched = content

	for handler in INJECT_HANDLERS:
		patched = patched.replace("void %s_Handler(void){"%(handler), "void %s_Handler(void){\nmy_%s();"%(handler, handler))

	with open(path, "w") as of:
		of.write(patched)

def inject_handlers(path: Path):
	root = path / "Core/Src"

	for child in root.iterdir():
		if any(child.match(ext) for ext in SOURCE_EXTS) and \
			not any(child.match(exclude) for exclude in CMAKE_EXCLUDE_SOURCES):
			inject_handlers_file(child)

def call_my_main(path: Path):
	main_c = path / "Core/Src/main.c"

	with open(main_c, "r") as _if:
		content = _if.read()

	patched = content.replace("/* USER CODE BEGIN PFP */", "/* USER CODE BEGIN PFP */\nvoid my_main();").replace("/* USER CODE BEGIN WHILE */", "/* USER CODE BEGIN WHILE */\nmy_main();")

	with open(main_c, "w") as of:
		of.write(patched)

def load_config(path: Path):
	if not path.exists():
		return
	if not path.is_file():
		print("Config '%s' is not a file."%(path))
		exit(1)
	
	with open(path, "rb") as cfg:
		config = toml.load(cfg)

	overlay_config(config)

def overlay_config(config: dict):
	global MX_DISCARD_PATHS
	global HEADER_EXTS
	global SOURCE_EXTS
	global CMAKE_TARGET
	global CMAKE_EXCLUDE_SOURCES
	global ASSUME_YES
	global INJECT_HANDLERS
	global NO_MY_MAIN

	if "discard" in config and config["discard"]:
		if not isinstance(config["discard"], list):
			print("discard must be a list of strings")
			exit(1)
		MX_DISCARD_PATHS += config["discard"]
	if "header" in config and config["header"]:
		if not isinstance(config["header"], list):
			print("header must be a list of strings")
			exit(1)
		HEADER_EXTS += config["header"]
	if "source" in config and config["source"]:
		if not isinstance(config["source"], list):
			print("source must be a list of strings")
			exit(1)
		SOURCE_EXTS += config["source"]
	if "target" in config and config["target"]:
		if not isinstance(config["target"], str):
			print("target must be a string")
			exit(1)
		CMAKE_TARGET = config["target"]
	if "exclude" in config and config["exclude"]:
		if not isinstance(config["exclude"], list):
			print("exclude must be a list of strings")
			exit(1)
		CMAKE_EXCLUDE_SOURCES = config["exclude"]
	if "assume-yes" in config and config["assume-yes"]:
		if not isinstance(config["assume-yes"], bool):
			print("assume-yes must be a bool")
			exit(1)
		ASSUME_YES |= config["assume-yes"]
	if "inject-handlers" in config and config["inject-handlers"]:
		if not isinstance(config["inject-handlers"], list):
			print("inject-handlers must be a list of strings")
			exit(1)
		INJECT_HANDLERS += config["inject-handlers"]
	if "no-my-main" in config and config["no-my-main"]:
		if not isinstance(config["no-my-main"], bool):
			print("no-my-main must be a bool")
			exit(1)
		NO_MY_MAIN |= config["no-my-main"]

def print_config():
	print("MX_DISCARD_PATHS", MX_DISCARD_PATHS)
	print("SOURCE_EXTS", SOURCE_EXTS)
	print("HEADER_EXTS", HEADER_EXTS)
	print("CMAKE_EXCLUDE_SOURCES", CMAKE_EXCLUDE_SOURCES)
	print("CMAKE_TARGET", CMAKE_TARGET)
	print("ASSUME_YES", ASSUME_YES)
	print("INJECT_HANDLERS", INJECT_HANDLERS)
	print("NO_MY_MAIN", NO_MY_MAIN)

def main2(ns: Namespace):
	if not ns.path.is_dir():
		print("'%s' is not a directory."%(ns.path))
		exit(1)

	load_config(ns.config)
	overlay_config(vars(ns))

	if ns.show_cfg:
		print_config()
		exit(0)

	mx_cleanup(ns.path)
	recurse_cmakelists(ns.path)
	root_cmakelists(ns.path)

	if not INJECT_HANDLERS:
		inject_handlers(ns.path)

	if not NO_MY_MAIN:
		call_my_main(ns.path)

def main():
	parser = ArgumentParser(prog="MX Convert")
	parser.add_argument("-c", "--config", action="store", type=Path, default="mx_convert.toml", help="The config file to read. If not specified, defaults to 'mx_convert.toml'.")
	parser.add_argument("-d", "--discard", action="append", help="Add a file to the discard list (delete if found).")
	parser.add_argument("-i", "--inject-handlers", action="store", help="Add a call to your own handler functions in MX generated code.")
	parser.add_argument("-H", "--header", action="append", help="Add a pattern to the headers list, single wildcards allowed. Includes *.h by default.")
	parser.add_argument("-n", "--no-my-main", action="store_true", help="Disable added call to my_main inside the generated main function.")
	parser.add_argument("-s", "--source", action="append", help="Add a pattern to the sources list, single wildcards allowed. Includes *.c and *.s by default.")
	parser.add_argument("-t", "--target", action="store", default="cubemx", help="The CMake target to generate references to. Defaults to 'cubemx'.")
	parser.add_argument("-x", "--exclude", action="append", help="Add a file to the source exclusion list, single wildcards allowed.")
	parser.add_argument("-y", "--assume-yes", action="store_true", help="Assume yes on all queries, useful for scripts.")
	parser.add_argument("--show-cfg", action="store_true", help="Show configs that would be used to run and exit successfully.")
	parser.add_argument("path", action="store", type=Path, help="The path to the directory to process.")

	namespace = parser.parse_args(argv[1:])

	main2(namespace)

if __name__ == "__main__":
	main()