#!/usr/bin/env python3

if __name__ == "__main__":
	print(":: C bytebeat generator: compiler unit")

from argparse import ArgumentParser
from os import environ, listdir, makedirs, name as os_name, \
               remove as delete_file, rmdir
from os.path import exists, join as path_join
from shlex import join as command_line_join, split as command_line_split
from shutil import which
from sys import stdin, stdout
from typing import Dict, Union
import re
import subprocess

# Definitions
BITS_PER_BYTE = 8
EXIT_FAILURE = 1
EXIT_SUCCESS = 0

# Paths
PATHS = {
	"src_dir": "src",
	"bin_dir": "bin",
	"template": "template.c",
	"substitute": "substituted.c",
	"output": "render_bytebeat",
	"fwrite_le_header": "fwrite_le.h",
	"fwrite_le": "fwrite_le.c",
	"include_directory": "include"
}

# Resolve paths
PATHS["template"] = path_join(PATHS["src_dir"], PATHS["template"])
PATHS["substitute"] = path_join(PATHS["bin_dir"], PATHS["substitute"])
PATHS["output"] = path_join(PATHS["bin_dir"], PATHS["output"])
PATHS["fwrite_le"] = path_join(PATHS["src_dir"], PATHS["fwrite_le"])

# Add `.` directory before all paths for compilation
for key in ["template", "substitute", "output", "fwrite_le",
            "include_directory"]:
	PATHS[key] = path_join(".", PATHS[key])

# Default parameters
DEFAULT_PARAMETERS = {
	"CC": "cc",
	"CFLAGS": "-Ofast -march=native -mtune=native -Wall -Wextra -Wpedantic "
	          "-pedantic -Wno-unused-variable -Wno-unused-but-set-variable "
	          "-Wno-dangling-else -Wno-parentheses -std=c99",
	"INPUT_FILE": PATHS["substitute"],
	"OUTPUT_FILE": PATHS["output"]
}

stdout_atty = hasattr(stdout, "isatty") and stdout.isatty()

def fetch(name: str):
	if from_env := environ.get(name):
		return from_env
	elif name != "CFLAGS_EXTRA":
		return DEFAULT_PARAMETERS[name]

def read_file(path: str) -> str:
	return open(path, "r", encoding="utf-8-sig").read()

def rewrite_file(path: str, content: str):
	return open(path, "w", encoding="utf-8").write(content)

def read_from_file_or_stdin(path: str) -> str:
	if path == "-":
		print("Reading from STDIN...", flush=True)
		return "\n".join(stdin)
	elif exists(path):
		return read_file(path)
	else:
		raise SystemExit(f"The specified file {path} doesn't exist")

def substitute_vars(replacements: Dict[str, Union[bool, str]], text: str,
                    verbose: bool) -> str:
	if verbose:
		print("Substituting values:")
	for placeholder, replacement in replacements.items():
		if isinstance(replacement, bool):
			replacement = preprocessor_bool(replacement)

		if verbose and placeholder != "bytebeat_contents":
			print(placeholder, ": ", replacement, sep="")
		text = text.replace(f"`{placeholder}`", str(replacement))
	if verbose:
		print()
	return text

def run_command(*command: list[str]) -> None:
	print(command_line_join(command), flush=True)
	if subprocess.run(command).returncode != EXIT_SUCCESS:
		raise SystemExit(EXIT_FAILURE)

def delete_empty_dir(path: str) -> None:
	if exists(path) and len(listdir(path)) == 0:
		rmdir(path)

preprocessor_bool = lambda value: "1" if value else "0"
C_str_repr = lambda s: '"' + s.replace("\\", "\\\\").replace(r'"', r'\"') + '"'

CC = fetch("CC")

CC_SEARCH_LIST = [
	"gcc",
	"clang",
	"tcc"
]
if os_name == "nt":
	CC_SEARCH_LIST = [
		"msc",
		*CC_SEARCH_LIST
	]

CFLAGS = fetch("CFLAGS")
INPUT_FILE = fetch("INPUT_FILE")
OUTPUT_FILE = fetch("OUTPUT_FILE")
if extra := fetch("CFLAGS_EXTRA"):
	CFLAGS += " " + extra

parameter_line_regex = re.compile(r"^\/\/ *RENDER PARAMETERS *: ", re.MULTILINE)

ALLOWED_ARGUMENTS_FROM_FILE = (
	"sample_rate",
	"final_sample_rate",
	"bit_depth",
	"signed",
	"channels",
	"no_return",
)

DEFAULT_ARGS = {
	"sample_rate": 8000,
	"final_sample_rate": 0,
	"bit_depth": 8,
	"signed": False,
	"channels": 1,
	"no_return": False
}

is_cmd_available = lambda cmd: which(cmd) is not None
is_cmd_unavailable = lambda cmd: which(cmd) is None

if __name__ == "__main__":
	parser = ArgumentParser(description=\
		"Substitutes supplied C (non-JavaScript!) bytebeat into the template, "
		"then attempts to compile the instance of the template. Accepts "
		"environmental variables `CC`, `CFLAGS`, `INPUT_FILE`, `OUTPUT_FILE`. "
		"`CFLAGS_EXTRA` can be used to add to default `CFLAGS`.")
	parser.add_argument("file", type=str,
		help="bytebeat formula file (use `-` to read from stdin)")
	parser.add_argument("-o", "--output", default="output.wav", type=str,
		help="specify output file path : default is `output.wav`")
	parser.add_argument("-r", "--sample-rate", default=None, type=int,
		help="sample rate (Hz)")
	parser.add_argument("-p", "--final-sample-rate", default=None, type=int,
		help="convert the output to a different sample rate (usually one that "
		"is set in the system, to improve sound quality) during generation "
		"(not just reinterpretation)")
	parser.add_argument("-b", "--bit-depth", default=None, type=int,
		help="bit depth")
	parser.add_argument("-s", "--signed", default=None, action="store_true",
		help="is signed?")
	parser.add_argument("-u", "--unsigned", default=None, action="store_true",
		help="is unsigned? (overrides the 'is signed' parameter)")
	parser.add_argument("-R", "--precalculate-ratio", default=False,
		action="store_true",
		help="precalculate sample ratio to speed up rendering (may produce "
		     "inaccurate results)")
	parser.add_argument("-m", "--faster-sample-ratio-math", default=False,
		action="store_true",
		help="faster sample ratio math (implies argument -R)")
	parser.add_argument("-f", "--floating-point", default=False,
		action="store_true", help="use floating point as the return type")
	parser.add_argument("-c", "--channels", default=None, type=int,
		help="amount of channels")
	parser.add_argument("-t", "--seconds", default=None, type=int,
		help="length in seconds (samples = sample rate * seconds) : "
		     "default = 30 seconds")
	parser.add_argument("-l", "--samples", default=None, type=int,
		help="length in samples (adds to `-t`; supports negative numbers) : "
		     "default = seconds + 0 samples")
	parser.add_argument("-S", "--skip-first", default=None, type=str,
		help="skip first `A` seconds and `B` samples: in format `As`, `B` or "
		     "`AsB` : default = 0")
	parser.add_argument("-k", "--repeat", default=0, type=int,
		help="how many times to repeat the bytebeat : "
		     "default = 0")
	parser.add_argument("-a", "--no-return", default=None, action="store_true",
		help="do not insert return statement before the code")
	parser.add_argument("-U", "--mode", default="sequential", type=str,
		help="mode of writing: `sequential` or `instant` (the latter is not "
		     "recommended, since the whole result would be stored in RAM)")
	parser.add_argument("-n", "--block-size", default=65536, type=int,
		help="sequential mode only: block size of each sequence, bytes")
	parser.add_argument("-q", "--silent", default=False, action="store_true",
		help="do not output anything during generation")
	parser.add_argument("-v", "--verbose", default=False, action="store_true",
		help="show progress during generation")
	parser.add_argument("-E", "--show-substituted-values", default=False,
		action="store_true", help="show substituted values")
	parser.add_argument("--color", default="auto", type=str,
		help="ANSI escape codes. Set to 'always' to enable them, 'none' to "
		     "disable. Default: 'auto'.")
	parser.add_argument("--keep-files", default=False, action="store_true",
		help="keep generated files: substituted source code of runtime unit "
		     "and the executable it will be compiled to.")
	args = parser.parse_args()

	bytebeat_contents = read_from_file_or_stdin(args.file).strip()

	if not bytebeat_contents:
		raise SystemExit("Empty file or STDIN")

	# - Parse arguments from file
	used_parameter_line = False
	for line in bytebeat_contents.splitlines():
		if (match := re.search(parameter_line_regex, line)) and \
		   not used_parameter_line:
			used_parameter_line = True

			parsed_parameters = line[match.start(0):].split(",")
			for parameter in parsed_parameters:
				kv = [x.strip() for x in parameter.split("=")]

				key = None
				value = None

				if len(kv) == 1:
					key, value = kv[0], True
				elif len(kv) == 2:
					key, value = kv[0], int(kv[1])
				else:
					break

				# Apply the argument only if it was not used by user yet and is
				# allowed to be set
				if (key not in args or getattr(args, key) is None) and \
				   key in ALLOWED_ARGUMENTS_FROM_FILE:
					setattr(args, key, value)

	# - Set default values
	for key, value in DEFAULT_ARGS.items():
		if getattr(args, key) is None:
			setattr(args, key, value)

	if args.unsigned is True:
		args.signed = False

	# - Compilation
	makedirs(PATHS["bin_dir"], exist_ok=True)

	if not args.no_return:  # Insert `return` statement
		# XXX: The bytebeat code is enclosed in parentheses to allow for the
		# use of commas as a comma operator, enabling more formulas to function.
		bytebeat_contents = f"return ({bytebeat_contents})"

	original_sample_rate = args.sample_rate
	final_sample_rate_code = ""
	if args.faster_sample_ratio_math:
		args.precalculate_ratio = True
	if args.precalculate_ratio and args.final_sample_rate != 0:
		if args.faster_sample_ratio_math:
			sample_rate_ratio = args.sample_rate / args.final_sample_rate
			final_sample_rate_code = f"time *= {sample_rate_ratio}L;"
		else:
			sample_rate_ratio = args.final_sample_rate / args.sample_rate
			final_sample_rate_code = f"time /= {sample_rate_ratio}L;"
		args.sample_rate = args.final_sample_rate

	final_sample_rate = \
		value if (value := args.final_sample_rate) != 0 \
		      else original_sample_rate
	samples = 0
	while True:
		no_seconds = args.seconds is None or args.seconds == 0
		no_samples = args.samples is None or args.samples == 0
		seconds_specified = not no_seconds
		samples_specified = not no_samples

		if seconds_specified and args.seconds < 0:
			raise SystemExit("CLI: Count of seconds can't be less than zero.")

		if no_seconds and samples_specified:
			samples = args.samples
		elif seconds_specified and samples_specified:
			samples = args.seconds * final_sample_rate + args.samples
		elif seconds_specified and no_samples:
			samples = args.seconds * final_sample_rate
		elif no_seconds and no_samples:
			args.seconds = 30  # default
			continue
		else:
			raise SystemExit("CLI: Incorrect seconds/samples length format.")
		break

	if samples <= 0:
		raise SystemExit("CLI: Count of samples should be greater than zero.")

	if args.mode != "sequential" and args.mode != "instant":
		raise SystemExit(f"Invalid mode '{args.mode}'")

	gen_length = args.channels * samples

	ansi_escape_codes_supported = args.color == "auto" and stdout_atty or \
	                              args.color == "always"

	actual_sample_rate = \
		value if (value := args.final_sample_rate) else args.sample_rate

	# - Parse the '--skip-first' argument
	if not args.skip_first is None:
		encountered_s = False
		for character in args.skip_first:
			if character.isdigit() or character == "s" and not encountered_s:
				if character == "s":
					encountered_s = True
			else:
				print(f"Invalid --skip-first format: `{args.skip_first}`")
				exit(1)

		skip_first = \
			[int(x) if x.isdigit() else 0 for x in args.skip_first.split("s")]

		skip_first_samples = 0
		if len(skip_first) == 1:
			skip_first += [0]

		skip_first_samples = skip_first[0] * actual_sample_rate + skip_first[1]
	else:
		skip_first_samples = 0

	length_formula = lambda channels, samples, n: channels * (samples + n)
	gen_length = length_formula(args.channels, samples, 0)
	loop_end = length_formula(args.channels, samples, skip_first_samples)

	rewrite_file(PATHS["substitute"], substitute_vars({
		"bytebeat_contents": bytebeat_contents,
		"output_file": C_str_repr(args.output),
		"sample_rate": actual_sample_rate,
		"original_sample_rate": original_sample_rate,
		"final_sample_rate_code": final_sample_rate_code,
		"bit_depth": args.bit_depth,
		"is_signed": args.signed,
		"precalculated_ratio": args.precalculate_ratio,
		"faster_sample_ratio_math": args.precalculate_ratio,
		"fp_return_type": args.floating_point,
		"channels": args.channels,
		"running_length": samples,
		"loop_end": loop_end,
		"loop_end_minus_1": loop_end - 1,
		"initial_time": skip_first_samples,
		"repeat_times": args.repeat,
		"wav_product": gen_length * (args.bit_depth // BITS_PER_BYTE),
		"gen_length": gen_length,
		"sequential_mode": args.mode == "sequential",
		"block_size": args.block_size,
		"silent_mode": args.silent,
		"verbose_mode": args.verbose and not args.silent,
		"fwrite_le": PATHS["fwrite_le_header"],
		"ansi_escape_codes_supported": ansi_escape_codes_supported
	}, read_file(PATHS["template"]), args.show_substituted_values))

	if is_cmd_unavailable(CC):
		print(f"Compiler {CC} is not available, searching:")

		still_unavailable = True
		for compiler in CC_SEARCH_LIST:
			print(f"* Trying CC={compiler}", end="")
			if is_cmd_available(compiler):
				print(": OK")
				CC = compiler
				still_unavailable = False
				break
			else:
				print()

		if still_unavailable:
			raise SystemExit("Could not find an available compiler. Please "
			                 "specify it by setting\nan environmental variable "
			                 "CC.")

	# Compile
	print("Compiling")

	run_command(
		CC,
		*command_line_split(CFLAGS),
		INPUT_FILE,
		PATHS["fwrite_le"],
		"-o", OUTPUT_FILE,
		"-I" + PATHS["include_directory"]
	)
	run_command(OUTPUT_FILE)

	if not args.keep_files:
		delete_file(PATHS["substitute"])
		delete_file(OUTPUT_FILE)
		delete_empty_dir(PATHS["bin_dir"])