from lexer import Lexer
from plasma_parser import Parser
from compiler import Compiler
from AST import Program
import json
import time
import sys
from argparse import ArgumentParser, Namespace

from llvmlite import ir
import llvmlite.binding as llvm
from ctypes import CFUNCTYPE, c_int, c_float


def parse_arguments() -> Namespace:
	arg_parser: ArgumentParser = ArgumentParser(
		"Plasma",
		description="The compiler for the Plasma programming language."
	)

	arg_parser.add_argument("file_path", type=str, help="path to the entry point of your plasma program (ex. \"main.pla\")")
	arg_parser.add_argument("--tokens", action="store_true", help="display the tokenized code")
	arg_parser.add_argument("--ast", default=None, help="export the generated AST to a JSON file")
	arg_parser.add_argument("--llvm", default=None, help="export the generated LLVM IR to a file")
	arg_parser.add_argument("--silent", action="store_true", help="don't print anything!!!!")

	return arg_parser.parse_args()


if __name__ == "__main__":
	args = parse_arguments()
	if args.silent:
		sys.stdout = None

	with open(args.file_path) as f:
		code: str = f.read()

	if args.tokens:
		debug_lex: Lexer = Lexer(source=code)
		while debug_lex.current_char is not None:
			print(debug_lex.next_token())

	l: Lexer = Lexer(source=code)
	p: Parser = Parser(lexer=l)

	parse_st: float = time.time()
	program: Program = p.parse_program()
	parse_et: float = time.time()
	#print(f"Parsed in {round((parse_et - parse_st) * 1000, 6)} ms.")

	if len(p.errors) > 0:
		for err in p.errors:
			print(err)
		exit(1)

	if args.ast:
		print("===== PARSER DEBUG =====")
		#program: Program = p.parse_program()

		with open(args.ast, "w") as f:
			json.dump(program.json(), f, indent=4)

		print(f"Wrote AST to \"{args.ast}\" successfully.")

	c: Compiler = Compiler()

	compiler_st: float = time.time()
	c.compile(program)
	compiler_et: float = time.time()
	#print(f"Compiled in {round((compiler_et - compiler_st) * 1000, 6)} ms.")

	module: ir.Module = c.module
	module.triple = llvm.get_default_triple()

	if args.llvm:
		with open(args.llvm, "w") as f:
			f.write(str(module))

	llvm.initialize_native_target()
	llvm.initialize_native_asmparser()
	llvm.initialize_native_asmprinter()

	try:
		llvm_ir_parsed = llvm.parse_assembly(str(module))
		llvm_ir_parsed.verify()
	except Exception as e:
		print(e)
		raise

	target_machine = llvm.Target.from_default_triple().create_target_machine()

	engine = llvm.create_mcjit_compiler(llvm_ir_parsed, target_machine)
	engine.finalize_object()

	entry = engine.get_function_address("main")
	cfunc = CFUNCTYPE(c_int)(entry)

	st = time.time()

	result = cfunc()

	et = time.time()

	#print(f"\n\nProgram returned: {result}\n=== Executed in {round((et - st) * 1000, 6)} ms. ===")
	exit(result)