#load "Lexing.jai";
#load "Error.jai";
#load "Parsing.jai";
#load "Semantic_Analysis.jai";
#load "Codegen.jai";


add_define :: (env : *Environment, key : string) {
	for define : env.defines {
		if define == key {
			return;
		}
	}

	array_add(*env.defines, key);
}

remove_define :: (env : *Environment, key : string) {
	for define : env.defines {
		if define == key {
			env.defines[it_index] = env.defines[env.defines.count - 1];
		}
	}
}

Environment :: struct {
	defines : [..]string;
}

Shader_Compiler :: struct {
	environment : Environment;
}

Field_Kind :: enum {
	Int       :: 0;
	Half      :: 1;
	Float     :: 2;
	Double    :: 3;
	Texture2D :: 8;
	Sampler   :: 9;

	Function;
	Struct;
	Array;
}

Field_Type :: struct {
	kind : Field_Kind;

	name : string; //@Note(niels): for structs

	children : [..]Field;
}

Hint_Kind :: enum {
	None;
	
	Position;
	UV;
	Target;

	Custom;
}

Field_Hint :: struct {
	kind : Hint_Kind;

	target_index : int;
	custom_hint_name : string;
}

Field :: struct {
	name : string;

	type : Field_Type;
	resource_index : u32;
	hints : [..]Field_Hint;
}

Entry_Point :: struct {
	name : string;

	function_input : [..]Field;
	return_value : Field;
}

Shader_Variant :: struct {
	text : string;

	vertex_entry_point : struct {
		name : string;

		input : [..]Field;
	}

	pixel_entry_point : struct {
		name : string;

		return_value : Field;
	}
}

Property_Field :: struct {
	base_field : Field;

	// @Incomplete(nb): Editor information, min max, etc.
	//                  This should also be compiled out for ship
}

Properties :: struct {
	fields : [..]Property_Field;

	buffer_index : u32;
}

Constant_Buffer :: struct {
	register : int;

	name : string;

	fields : Static_Array(Property_Field, 16);

	buffer_index : u32;
}

Shader_Variant_Collection :: struct {
	properties : Properties;

	max_constant_buffers :: 16;
	cbuffers : Static_Array(Constant_Buffer, max_constant_buffers);

	variants : [..]Shader_Variant;
}

Input_File :: struct {
	source : string;
	path   : string;
}

Token_Stream :: struct {
	tokens : [..]Token;
}

Compiled_File :: struct {
	file : Input_File;
	tokens : Token_Stream;
	ast_root  : *AST_Node;
	ast_nodes : [..]AST_Node;

	codegen_result_text : string;

	semantic_check_result : Semantic_Check_Result;

	vertex_entry_point : struct {
		name : string;

		input : [..]Field;
	}

	pixel_entry_point : struct {
		name : string;

		return_value : Field;
	}

	properties : Properties;

	max_constant_buffers :: 16;
	cbuffers : Static_Array(Constant_Buffer, max_constant_buffers);
}

Compile_Result :: struct {	
	files : [..]Compiled_File;

	had_error : bool;
	messages : [..]Compiler_Message;

	allocator : Allocator;
	arena     : Arena;
}

//@Incomplete(niels): need to consider allocation
add_file :: (result : *Compile_Result, path : string) {
	file_string, ok := read_entire_file(path);

	if !ok {
		// record_error(.File_Load_Failed, "Unable to load file: %", path);
		return;
	}

	input_file : Input_File;

	input_file.source = file_string;
	input_file.path   = path;

	compiled_file : Compiled_File;
	compiled_file.file = input_file;

	array_add(*result.files, compiled_file);
}

// @Incomplete(nb): Will we ever even use this?
from_file :: (path : string) -> Compile_Result {
	arr : [1]string;
	arr[0] = path;
	return from_files(arr);
}

from_files :: (paths : []string) -> Compile_Result {
	result : Compile_Result;
	for path : paths {
		add_file(*result, path);
	}

	return result;
}

Compilation_Result :: struct {
	messages : [..]Compiler_Message;

	had_error : bool;

	collection : Shader_Variant_Collection;
}

pretty_print_field :: (field : *Field) -> string {
	builder : String_Builder;
	init_string_builder(*builder,, temp);

	pretty_print_field(*builder, field);

	return builder_to_string(*builder);
}

Min_Field_Name :: 10;

pretty_print_field :: (builder : *String_Builder, field : *Field) {
	if field.name.count > 0 {
		print_to_builder(builder, "% ", field.name);
		append(builder, "- ");
	} else {
		append(builder, "return - ");
	}

	type := field.type;

	if type.kind == {
		case .Int; {
			append(builder, "int");
		}
		case .Half; {
			append(builder, "half");
		}
		case .Float; {
			append(builder, "float");
		}
		case .Double; {
			append(builder, "double");
		}
		case .Texture2D; {
			append(builder, "texture2D");
		}
		case .Sampler; {
			append(builder, "sampler");
		}
		case .Struct; {
			print_to_builder(builder, "struct : % {", type.name);

			for *child : type.children {
				pretty_print_field(builder, child);
				if it_index < type.children.count - 1 {
					append(builder, " ");
				}
			}

			append(builder, "} ");
		}
		case .Array; {

		}
	}

	for hint : field.hints {
		if hint.kind == {
			case .Position; {
				append(builder, "(@position)");
			}
			case .Target; {
				print_to_builder(builder, "(@target%)", hint.target_index);
			}
			case .Custom; {
				print_to_builder(builder, "(@%)", hint.custom_hint_name);
			}
		}

		if it_index != field.hints.count - 1 {
			append(builder, ", ");
		}
	}
}

type_variable_to_field :: (checker : *Semantic_Checker, variable : Type_Variable_Handle) -> Field {
	return type_variable_to_field(checker, from_handle(checker, variable));
}

type_variable_to_field :: (type_variables : []Type_Variable, scope_stack : Scope_Stack, variable : *Type_Variable) -> Field {
	field : Field;

	field.name = variable.name;

	type : Field_Type;

	if variable.type == {
		case .Int; {
			type.kind = Field_Kind.Int;
		}
		case .Half; {
			type.kind = Field_Kind.Half;
		}
		case .Float; {
			type.kind = Field_Kind.Float;
		}
		case .Double; {
			type.kind = Field_Kind.Double;
		}
		case .Texture2D; {
			type.kind = Field_Kind.Texture2D;

			field.resource_index = variable.resource_index;
		}
		case .Sampler; {
			type.kind = Field_Kind.Sampler;

			field.resource_index = variable.resource_index;
		}
		case .Struct; {
			type.kind = Field_Kind.Struct;

			find_result := find_symbol(scope_stack, variable.typename, xx 1);
			assert(find_result != null, "Internal compiler error\n");

			type_var := from_handle(type_variables, find_result.type_variable);

			for i : 0..type_var.children.count - 1 {
				child := type_var.children[i];
				child_field := type_variable_to_field(type_variables, scope_stack, child);
				array_add(*type.children, child_field);
			}

			type.name = variable.typename;
		}
	}

	for hint : variable.source_node.hint_tokens {
		field_hint : Field_Hint;

		if hint.ident_value == "position" {
			// @Incomplete(nb): Should be a lookup table somewhere
			field_hint.kind = .Position;
		} else if hint.ident_value == "uv" {
			field_hint.kind = .UV;
		} else if starts_with(hint.ident_value, "target") {
			// @Incomplete(nb): Should be a lookup table somewhere
			index_str : string;
			index_str.data = *hint.ident_value.data[7];
			index_str.count = 1;
			
			result, ok, remainder := string_to_int(index_str);
			if ok {
				field_hint.target_index = result;
			}
			field_hint.kind = .Target;
		} else {
			// @Incomplete(nb): custo hints
		}
		array_add(*field.hints, field_hint);
	}

	field.type = type;

	return field;
}

type_variable_to_field :: (type_variables : []Type_Variable, scope_stack : Scope_Stack, variable : Type_Variable_Handle) -> Field {
	return type_variable_to_field(type_variables, scope_stack, from_handle(type_variables, variable));
}

type_variable_to_field :: (checker : *Semantic_Checker, variable : *Type_Variable) -> Field {
	return type_variable_to_field(checker.result.type_variables, checker.result.scope_stack, variable);
}

compile_file :: (compiler : *Shader_Compiler, paths : []string) -> Compile_Result {
	result : Compile_Result;

	for path : paths {
		add_file(*result, path);
	}

	lex(*result);
	parse(*result);
	check(*result);
	codegen(*result);

	for *file : result.files {
		check_result := file.semantic_check_result;
		if check_result.vertex_entry_point {
			file.vertex_entry_point.name = check_result.vertex_entry_point.name;

			type_variable := from_handle(check_result.type_variables, check_result.vertex_entry_point.type_variable);
			assert(type_variable.type == .Function);

			node := type_variable.source_node;
			if node.children.count > 0 {
				if node.children[0].kind == .FieldList {
					field_list := node.children[0];
					for child : field_list.children {
						tv := from_handle(check_result.type_variables, child.type_variable);
						field := type_variable_to_field(check_result.type_variables, check_result.scope_stack, tv);
						array_add(*file.vertex_entry_point.input, field);
					}
				}
			}
		}

		for buffer_variable : to_array(*check_result.constant_buffers) {
			variable := from_handle(check_result.type_variables, buffer_variable);

			cb := array_add(*file.cbuffers);

			for i : 0..variable.children.count - 1 {
				child := variable.children[i];
				field : Property_Field;
				field.base_field = type_variable_to_field(check_result.type_variables, check_result.scope_stack, from_handle(check_result.type_variables, child));
				array_add(*cb.fields, field);
			}

			cb.buffer_index = variable.resource_index;
		}

		find_result := find_symbol(*check_result.scope_stack, check_result.property_name, xx 1);
		if find_result {
			property_variable := from_handle(check_result.type_variables, find_result.type_variable);

			for i : 0..property_variable.children.count - 1 {
				child := property_variable.children[i];
				field := type_variable_to_field(check_result.type_variables, check_result.scope_stack, from_handle(check_result.type_variables, child));
				prop_field : Property_Field;
				prop_field.base_field = field;
				array_add(*file.properties.fields, prop_field);
			}
			file.properties.buffer_index = property_variable.resource_index;
		}

		if check_result.pixel_entry_point {
			file.pixel_entry_point.name = check_result.pixel_entry_point.name;

			type_variable := from_handle(check_result.type_variables, check_result.pixel_entry_point.type_variable);
			assert(type_variable.type == .Function);

			field := type_variable_to_field(check_result.type_variables, check_result.scope_stack, type_variable.return_type_variable);
			for hint : type_variable.source_node.hint_tokens {
				field_hint : Field_Hint;

				if hint.ident_value == "position" {
					// @Incomplete(nb): Should be a lookup table somewhere
					field_hint.kind = .Position;
				} else if starts_with(hint.ident_value, "target") {
					// @Incomplete(nb): Should be a lookup table somewhere
					index_str : string;
					index_str.data = *hint.ident_value.data[7];
					index_str.count = 1;
					
					result, ok, remainder := string_to_int(index_str);
					if ok {
						field_hint.target_index = result;
					}
					field_hint.kind = .Target;
				} else {
					// @Incomplete(nb): custom hints
				}
				array_add(*field.hints, field_hint);
			}

			file.pixel_entry_point.return_value = field;
		}
	}
	
	return result;
}

compile_file :: (compiler : *Shader_Compiler, path : string) -> Compilation_Result {
	result : Compilation_Result;

	lexer : Lexer;
	
	init_lexer_from_file(*lexer, path);
	if lexer.result.had_error {
		copy_messages(lexer.result.messages, *result.messages);
		result.had_error = true;
		return result;
	}
	
	lex_result := lex(*lexer,, *temp);
	if lex_result.had_error {
		copy_messages(lex_result.messages, *result.messages);
		result.had_error = true;
		return result;
	}

	parse_state : Parse_State;
	init_parse_state(*parse_state, lex_result.tokens, lexer.path);

	parse_result := parse(*parse_state);
	if parse_result.had_error {
		copy_messages(parse_result.messages, *result.messages);
		result.had_error = true;
		return result;
	}

	checker : Semantic_Checker;
	init_semantic_checker(*checker, parse_result.root, path);

	check_result := check(*checker);
	if check_result.had_error {
		copy_messages(check_result.messages, *result.messages);
		result.had_error = true;
		return result;
	}

	state : Codegen_State;
	init_codegen_state(*state, parse_result.root, check_result, .HLSL);

	result_text : string;

	codegen_result := codegen(*state);
	if codegen_result.had_error {
		copy_messages(codegen_result.messages, *result.messages);
		result.had_error = true;
		
		return result;
	}

	// @Incomplete(nb): Assumes only a single variant for now

	variant : Shader_Variant;
	variant.text = codegen_result.result_text;

	if checker.result.vertex_entry_point {
		variant.vertex_entry_point.name = checker.result.vertex_entry_point.name;

		type_variable := from_handle(*checker, checker.result.vertex_entry_point.type_variable);
		assert(type_variable.type == .Function);

		node := type_variable.source_node;
		if node.children.count > 0 {
			if node.children[0].kind == .FieldList {
				field_list := node.children[0];
				for child : field_list.children {
					tv := from_handle(*checker, child.type_variable);
					field := type_variable_to_field(*checker, tv);
					array_add(*variant.vertex_entry_point.input, field);
				}
			}
		}
	}

	for buffer_variable : to_array(*check_result.constant_buffers) {
		variable := from_handle(check_result.type_variables, buffer_variable);

		cb := array_add(*result.collection.cbuffers);

		for i : 0..variable.children.count - 1 {
			child := variable.children[i];
			field : Property_Field;
			field.base_field = type_variable_to_field(*checker, from_handle(*checker, child));;
			array_add(*cb.fields, field);
		}

		cb.buffer_index = variable.resource_index;
	}

	find_result := find_symbol(*check_result.scope_stack, check_result.property_name, xx 1);
	if find_result {
		property_variable := from_handle(check_result.type_variables, find_result.type_variable);

		for i : 0..property_variable.children.count - 1 {
			child := property_variable.children[i];
			field := type_variable_to_field(*checker, from_handle(*checker, child));
			prop_field : Property_Field;
			prop_field.base_field = field;
			array_add(*result.collection.properties.fields, prop_field);
		}
		result.collection.properties.buffer_index = property_variable.resource_index;
	}

	if checker.result.pixel_entry_point {
		variant.pixel_entry_point.name = checker.result.pixel_entry_point.name;

		type_variable := from_handle(*checker, checker.result.pixel_entry_point.type_variable);
		assert(type_variable.type == .Function);

		field := type_variable_to_field(*checker, type_variable.return_type_variable);
		for hint : type_variable.source_node.hint_tokens {
			field_hint : Field_Hint;

			if hint.ident_value == "position" {
				// @Incomplete(nb): Should be a lookup table somewhere
				field_hint.kind = .Position;
			} else if starts_with(hint.ident_value, "target") {
				// @Incomplete(nb): Should be a lookup table somewhere
				index_str : string;
				index_str.data = *hint.ident_value.data[7];
				index_str.count = 1;
				
				result, ok, remainder := string_to_int(index_str);
				if ok {
					field_hint.target_index = result;
				}
				field_hint.kind = .Target;
			} else {
				// @Incomplete(nb): custo hints
			}
			array_add(*field.hints, field_hint);
		}

		variant.pixel_entry_point.return_value = field;
	}

	array_add(*result.collection.variants, variant);

	return result;
}