import "!hashmap.pry";

import "llvm.pry";

let Variable = struct {
	value: LLVMValueRef,
	type: LLVMTypeRef,
	node: *Node,
	node_type: *Node,
	stack_level: *i64,
};

let Scope = struct {
	variables: *HashMap,
};

let Environment = struct {
	scope_stack: **Scope,
	scope_stack_len: i64,
	arena: *arena,
};

let environment_create_scope = (e: *Environment) => void {
	let scope = cast(*Scope, arena_alloc((*e).arena, sizeof(Scope)));
	(*scope).variables = hashmap_init(16, (*e).arena);
	(*((*e).scope_stack + cast(**Scope, (*e).scope_stack_len))) = scope;
	(*e).scope_stack_len = (*e).scope_stack_len + 1;

	return;
};

let environment_drop_scope = (e: *Environment) => void {
	(*e).scope_stack_len = (*e).scope_stack_len - 1;

	return;
};

/* TODO: */
let environment_add_variale = (e: *Environment, name: *i8, variable: *Variable) => void {
	return;
};

/* TODO: */
let environment_set_variale = (e: *Environment, name: *i8, varbiable: *Variable) => void {
	return;
};

let environment_init = (alloc: *arena) => *Environment {
	let e = cast(*Environment, arena_alloc(alloc, sizeof(Environment)));
	(*e).scope_stack = cast(**Scope, arena_alloc(alloc, sizeof(*Scope) * 20));
	(*e).scope_stack_len = 0;
	(*e).arena = alloc;

	environment_create_scope(e);

	return e;
};

let codegen = struct {
	llvm_module: LLVMModuleRef,
	llvm_context: LLVMContextRef,
	builder: LLVMBuilderRef,
	arena: *arena,
	environment: *Environment,

	current_function: LLVMValueRef,
	current_function_retur_type: *Node,
};

let codegen_init = (alloc: *arena) => *codegen {
	LLVMInitializeAllTargetInfos();
	LLVMInitializeAllTargetMCs();
	LLVMInitializeAllTargets();
	LLVMInitializeAllAsmPrinters();
	LLVMInitializeAllAsmParsers();

	let module = LLVMModuleCreateWithName("module");
        let context = LLVMGetGlobalContext();
        let builder = LLVMCreateBuilder();
	
	let c = cast(*codegen, arena_alloc(alloc, sizeof(codegen)));
	
	(*c).llvm_module = module;
	(*c).llvm_context = context;
	(*c).builder = builder;
	(*c).arena = alloc;
	(*c).environment = environment_init(alloc);

	return c;
};

let create_node = (c: *codegen, n: Node) => *Node {
	let res = cast(*Node, arena_alloc((*c).arena, sizeof(Node)));
	*res = n;
	return res;
};

let codegen_create_variable = (c: *codegen, variable: Variable) => *Variable {
	let v = cast(*Variable, arena_alloc((*c).arena, sizeof(Variable)));
	*v = variable;
	return v;
};

let codegen_get_llvm_type = (c: *codegen, node: *Node) => *LLVMTypeRef {
	assert((*node).type >= NODE_TYPE_SIMPLE_TYPE);
	assert((*node).type <= NODE_TYPE_STRUCT_TYPE);
	
	if (*node).type == NODE_TYPE_SIMPLE_TYPE {
		let simple_type = *cast(*NODE_TYPE_SIMPLE_TYPE_DATA, (*node).data);

		if strcmp(simple_type.name, "i8") {
			let r = cast(*LLVMTypeRef, arena_alloc((*c).arena, sizeof(LLVMTypeRef)));
			*r = LLVMInt8Type();
			return r;
		};
		
		if strcmp(simple_type.name, "i64") {
			let r = cast(*LLVMTypeRef, arena_alloc((*c).arena, sizeof(LLVMTypeRef)));
			*r = LLVMInt64Type();
			return r;
		};
		
		if strcmp(simple_type.name, "void") {
			let r = cast(*LLVMTypeRef, arena_alloc((*c).arena, sizeof(LLVMTypeRef)));
			*r = LLVMVoidType();
			return r;
		};
		
		if strcmp(simple_type.name, "varargs") { /* Hack for varargs (only used for printf) */
			let r = cast(*LLVMTypeRef, arena_alloc((*c).arena, sizeof(LLVMTypeRef)));
			*r = LLVMPointerType(LLVMInt64Type(), 0);
			return r;
		};

		printf("NO SIMPLE TYPE %s!\n", simple_type.name);
		assert(false);
	};
	
	if (*node).type == NODE_TYPE_FUNCTION_TYPE {
		let function_type = *cast(*NODE_TYPE_FUNCTION_TYPE_DATA, (*node).data);
		let retur_type = codegen_get_llvm_type(c, function_type.retur_type);
		if retur_type == cast(*LLVMTypeRef, null) {
			return cast(*LLVMTypeRef, null);
		};
		let paramtypes = cast(*LLVMTypeRef, arena_alloc((*c).arena, sizeof(LLVMTypeRef) * 20));
		let paramtypes_len = 0;
		let is_varargs = 0;

		let i = 0;
		while i < function_type.parameters_len {
			let param = *(function_type.parameters + cast(**Node, i));
			if (*param).type == NODE_TYPE_SIMPLE_TYPE {
				let simple_type = *cast(*NODE_TYPE_SIMPLE_TYPE_DATA, (*param).data);
				if strcmp(simple_type.name, "varargs") {
					is_varargs = 1;
					i = i + 1;
					continue;
				};
			};
			let typ = codegen_get_llvm_type(c, param);
			if typ == cast(*LLVMTypeRef, null) {
				return cast(*LLVMTypeRef, null);
			};
			if (*param).type == NODE_TYPE_FUNCTION_TYPE {
				*typ = LLVMPointerType(*typ, 0);
			};

			(*(paramtypes + cast(*LLVMTypeRef, paramtypes_len))) = *typ;
			paramtypes_len = paramtypes_len + 1;

			i = i + 1;
		};
		let function_type = LLVMFunctionType(*retur_type, paramtypes, paramtypes_len, is_varargs);
		let r = cast(*LLVMTypeRef, arena_alloc((*c).arena, sizeof(LLVMTypeRef)));
		*r = function_type;
		printf("FUNCTION\n");
		return r;
	};

	if (*node).type == NODE_TYPE_POINTER_TYPE {
		let pointer_type = *cast(*NODE_TYPE_POINTER_TYPE_DATA, (*node).data);
		let inner_type = codegen_get_llvm_type(c, pointer_type.type);
		if inner_type == cast(*LLVMTypeRef, null) {
			return cast(*LLVMTypeRef, null);
		};
		let r = cast(*LLVMTypeRef, arena_alloc((*c).arena, sizeof(LLVMTypeRef)));
		*r = LLVMPointerType(*inner_type, 0);
		return r;
	};

	printf("NO TYPEEE BOI %d\n", (*node).type);
	assert(false);

	return cast(*LLVMTypeRef, null);
};

let codegen_generate_literal = (c: *codegen, literal_val: LLVMValueRef, name: *i8, node: *Node, node_type: *Node) => *Variable {
	/* TODO: Global */
	let v = Variable{};
	v.value = literal_val;
	v.type = cast(LLVMTypeRef, null);
	v.stack_level = cast(*i64, null);
	v.node = node;
	v.node_type = node_type;
	return codegen_create_variable(c, v);
};

extern codegen_generate_statement = (*codegen, *Node) => i64;

let codegen_generate_expression_value = (c: *codegen, expression: *Node, name: *i8) => *Variable {
	printf("NAME: %s\n", name);
	if ((*expression).type == NODE_PRIMARY_EXPRESSION_NUMBER) {
		printf("THIS\n");
		let n = (*cast(*NODE_PRIMARY_EXPRESSION_NUMBER_DATA, (*expression).data)).value;
		
		printf("X: %d\n", n);

		let node_type = Node{};
		node_type.type = NODE_TYPE_SIMPLE_TYPE;

		let d = cast(*NODE_TYPE_SIMPLE_TYPE_DATA, arena_alloc((*c).arena, sizeof(NODE_TYPE_SIMPLE_TYPE_DATA)));
		(*d).name = "i64";
		(*d).underlying_type = cast(*Node, null);
		node_type.data = cast(*void, d);
		
		return codegen_generate_literal(c, LLVMConstInt(LLVMInt64Type(), n, 0), name, expression, create_node(c, node_type));
	};
	
	if ((*expression).type == NODE_FUNCTION_DEFINITION) {
		printf("THIS2\n");
	/* TODO: IMPLEMENT */
		printf("ASS %d\n", (*expression).type);

		let builder_pos = LLVMGetInsertBlock((*c).builder);

		let llvm_param_types = cast(*LLVMTypeRef, arena_alloc((*c).arena, sizeof(LLVMTypeRef) * 20));
		let param_types = cast(**Node, arena_alloc((*c).arena, sizeof(*Node) * 20));
		/* TODO: VARARGS */

		
		let function_definition = *cast(*NODE_FUNCTION_DEFINITION_DATA, (*expression).data);

		let i = 0;
		/* TODO */
		while i < function_definition.parameters_len {
			i = i + 1;
		};

		let retur_type = codegen_get_llvm_type(c, function_definition.retur_type);
		if retur_type == cast(*LLVMTypeRef, null) {
			printf("LEL2\n");
			return cast(*Variable, null);
		};

		let function_type = LLVMFunctionType(*retur_type, llvm_param_types, i, 0);
		let function = LLVMAddFunction((*c).llvm_module, name, function_type);
                let function_entry = LLVMAppendBasicBlock(function, "entrypoint");
                LLVMPositionBuilderAtEnd((*c).builder, function_entry);

		/* TODO: SCOPE */
		environment_create_scope((*c).environment);
		let last_function = (*c).current_function;
		(*c).current_function = function;
		let last_function_retur_type = (*c).current_function_retur_type;
		(*c).current_function_retur_type = function_definition.retur_type;

		/* TODO: Defer */

		let d = cast(*NODE_TYPE_FUNCTION_TYPE_DATA, arena_alloc((*c).arena, sizeof(NODE_TYPE_FUNCTION_TYPE_DATA)));
		(*d).parameters = param_types;
		(*d).parameters_len = i;
		(*d).retur_type = function_definition.retur_type;
		let n = Node{};
		let node_type = create_node(c, n);
		(*node_type).type = NODE_TYPE_FUNCTION_TYPE;
		(*node_type).data = cast(*void, d);

		/* TODO: Recurisve functions */

		let params = cast(*LLVMValueRef, arena_alloc((*c).arena, sizeof(LLVMValueRef) * function_definition.parameters_len));
		LLVMGetParams(function, params);


		/* TODO */
		let parameters_index = 0;
		while parameters_index < function_definition.parameters_len {
			parameters_index = parameters_index + 1;
		};

		i = 0;
		while i < function_definition.statements_len {
			let stmt = *(function_definition.statements + cast(**Node, i));

			let res = codegen_generate_statement(c, stmt);
			if res != 0 {
				printf("LEL\n");
				return cast(*Variable, null);
			};

			i = i + 1;
		};

		LLVMPositionBuilderAtEnd((*c).builder, builder_pos);
	
		let v = Variable{};
		v.value = function;
		v.type = cast(LLVMTypeRef, null);
		v.stack_level = cast(*i64, null);
		v.node = expression;
		v.node_type = node_type;
		return codegen_create_variable(c, v);
	};
	
	if ((*expression).type == NODE_TYPE_FUNCTION_TYPE) {
		let e = *((*c).environment);
		printf("LLEL\n");
		assert(e.scope_stack_len == 1);
		printf("THIS3\n");

		/* TODO: Check if already exists */

		let function_type = codegen_get_llvm_type(c, expression);
		if function_type == cast(*LLVMTypeRef, null) {
			printf("LOL\n");
			return cast(*Variable, null);
		};
		let function = LLVMAddFunction((*c).llvm_module, name, *function_type);
		let v = Variable{};
		v.value = function;
		v.type = cast(LLVMTypeRef, null);
		v.stack_level = cast(*i64, null);
		v.node = expression;
		v.node_type = expression;
		printf("RET\n");
		return codegen_create_variable(c, v);
	};

	printf("ASSERT 1\n");
	assert(false);

	return cast(*Variable, null);
};

let codegen_generate_assignment_statement = (c: *codegen, stmt: *NODE_ASSIGNMENT_STATEMENT_DATA) => i64 {
	let lhs = *((*stmt).lhs);
	let prhs = (*stmt).rhs;

	if (lhs.type == NODE_PRIMARY_EXPRESSION_IDENTIFIER) {
		let identifier = (*cast(*NODE_PRIMARY_EXPRESSION_IDENTIFIER_DATA, lhs.data)).name;
		printf("XX %s\n", identifier);
		let variable = codegen_generate_expression_value(c, prhs, identifier);
		assert(variable != cast(*Variable, null));

		/* TODO IMPORTANT: Add variable to variables */

		return 0;
	};

	printf("ASSERT 2 %d\n", lhs.type);
	assert(false);
	return 0;
};

let codegen_generate_return_statement = (c: *codegen, stmt: *NODE_RETURN_STATEMENT_DATA) => i64 {
	printf("HELLO!\n");
	let expression = (*stmt).expression;

	if expression == cast(*Node, null) {
		printf("RETURN VOID!\n");
		LLVMBuildRetVoid((*c).builder);
		return 0;
	};

	let val = codegen_generate_expression_value(c, expression, cast(*i8, null));
	if val == cast(*Variable, null) {
		return 1;
	};

	/* TODO: Check type */

	LLVMBuildRet((*c).builder, (*val).value);

	printf("BYE!\n");

	return 0;
};

let codegen_generate_function_call_statement = (c: *codegen, stmt: *NODE_FUNCTION_CALL_STATEMENT_DATA) => i64 {
	printf("LELELEL\n");
	assert(false);

	return 0;
};

let codegen_generate_statement = (c: *codegen, statement: *Node) => i64 {
	let stmt = *statement;

	if stmt.type == NODE_ASSIGNMENT_STATEMENT {
		return codegen_generate_assignment_statement(c, cast(*NODE_ASSIGNMENT_STATEMENT_DATA, stmt.data));
	};

	if stmt.type == NODE_RETURN_STATEMENT {
		return codegen_generate_return_statement(c, cast(*NODE_RETURN_STATEMENT_DATA, stmt.data));
	};
	
	if stmt.type == NODE_FUNCTION_CALL_STATEMENT {
		return codegen_generate_function_call_statement(c, cast(*NODE_FUNCTION_CALL_STATEMENT_DATA, stmt.data));
	};

	printf("ASSERT 3 %d\n", stmt.type);
	assert(false);

	return 0;
};

let codegen_generate = (c: *codegen, ast: *Node) => i64 {
	assert((*ast).type == NODE_PROGRAM);

	let program = *cast(*NODE_PROGRAM_DATA, (*ast).data);

	let i = 0;
	while i < program.statements_len {
		let stmt = *(program.statements + cast(**Node, i));

		let res = codegen_generate_statement(c, stmt);
		if res != 0 {
			return 1;
		};

		i = i + 1;
	};

	return 0;
};

let codegen_compile = (c: *codegen) => i64 {
        /* Dump module */
	LLVMDumpModule((*c).llvm_module);

	/* Generate code */
        let triple = LLVMGetDefaultTargetTriple();
        let target_ref = cast(*LLVMTargetRef, arena_alloc((*c).arena, sizeof(*LLVMTargetRef)));
        let message = cast(**i8, null);
        let result = LLVMGetTargetFromTriple(triple, target_ref, message);
        if result != 0 {
            printf("Target output: %s\n", *message);
            LLVMDisposeMessage(*message);
        };
        let target_machine = LLVMCreateTargetMachine(
            *target_ref,
            triple,
            "",
            "",
            LLVMCodeGenLevelDefault,
            LLVMRelocDefault,
            LLVMCodeModelDefault,
        );
        LLVMDisposeMessage(triple);
        result = LLVMVerifyModule((*c).llvm_module, LLVMAbortProcessAction, message);
        if result != 0 {
            printf("Verification output: %s\n", *message);
            LLVMDisposeMessage(*message);
        };

        /* Generate the object file */
        let filename = "bootstrap_output.o";
        LLVMTargetMachineEmitToFile(
            target_machine,
            (*c).llvm_module,
            filename,
            LLVMObjectFile,
            cast(**i8, null),
        );
        LLVMDisposeTargetMachine(target_machine);
        printf("Object file generated: %s\n", filename);

	return 0;
};

let codegen_deinit = (c: *codegen) => void {
        LLVMDisposeModule((*c).llvm_module);
        LLVMShutdown();
        LLVMDisposeBuilder((*c).builder);
	return;
};