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

#import "File_Utilities";

/* TODO
- [x] Remove builtin stringbuilding and replace it with ad-hoc string building when error reporting. In that case we are already building a string anyway, so we can just pass in the string builder
- [ ] Support structured buffers (ro, rw, w)
- [ ] Support mesh and amplification shaders
- [ ] Support compute shaders
- [x] Support #if at top level
- [x] Support #if at block level
- [x] Remove properties block and just use hinted constant buffers instead
     ```
     props :: constant_buffer @properties {
          [...]
     }
     ```
- [ ] while loops
- [ ] for-each loops
- [ ] add parameters to hints (meta properties, resource binding indices if needed)
- [ ] consider @entry(stage) syntax instead of the forced keyword
*/

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;
}

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;
	Output_Position;

	Custom;
}

Hint_Names :: #run -> [(cast(int)Hint_Kind.Target) + 1]string {
	names : [(cast(int)Hint_Kind.Target) + 1]string;
	names[Hint_Kind.Position] = "position";
	names[Hint_Kind.UV]       = "uv";
	names[Hint_Kind.Target]   = "target";

	return names;
}

lookup_hint :: (name : string) -> Hint_Kind {
	if name == "position" {
		return Hint_Kind.Position;		
	} else if name == "uv" {
		return Hint_Kind.UV;
	} else if starts_with(name, "target") {
		return Hint_Kind.Target;
	} else if name == "outposition" {
		return Hint_Kind.Output_Position;
	}
	return .None;
}

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;
}

Constant_Buffer :: struct {
	name : string;

	fields : Static_Array(Field, 16);

	hints : [..]Field_Hint;

	buffer_index : u32;
}

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

Compiler_Context :: struct {
	file : Input_File;

	environment : Environment;
	
	tokens : [..]Token;;
	root  : *AST_Node;
	nodes : [..]AST_Node;

	codegen_result_text : string;

	constant_buffers : Static_Array(Type_Variable_Handle, 16);
	buffers : Static_Array(Type_Variable_Handle, 16);

	scope_stack : Scope_Stack;
	type_variables : [..]Type_Variable;

	vertex_entry_point : struct {
		node : *AST_Node;
		name : string;
		input : [..]Field;
	}

	pixel_entry_point : struct {
		node : *AST_Node;
		name : string;
		return_value : Field;
	}

	max_constant_buffers :: 16;

	cbuffers : Static_Array(Constant_Buffer, max_constant_buffers);

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

#add_context scratch_allocators : [2]Allocator;
#add_context scratch_id : int = 0;

init_context_allocators :: () {
	if get_arena(context.scratch_allocators[0]) == null {
		context.scratch_allocators[0] = make_arena(Megabytes(128));
		context.scratch_allocators[1] = make_arena(Megabytes(128));
	}
}

clear_context_allocators :: () {
	if get_arena(context.scratch_allocators[0]) != null {
		clear(context.scratch_allocators[0]);
		clear(context.scratch_allocators[1]);
	}
}

get_scratch :: (conflict : Allocator = .{}) -> Scratch {
	arena := cast(*Arena)conflict.data;
	if arena == get_arena(context.scratch_allocators[0]) || context.scratch_id == 0 {
		context.scratch_id = 1;
		return scratch_begin(*context.scratch_allocators[1]);
	}
	context.scratch_id = 0;
	return scratch_begin(*context.scratch_allocators[0]);
}

record_error :: (result : *Compiler_Context, format : string, args : .. Any) {
	error : Compiler_Message;
	error.message_kind = .Error;
	error.message      = sprint(format, args);
	
	array_add(*result.messages, error);
}

make_file :: (result : *Compiler_Context, path : string) -> Input_File {
	if !file_exists(path) {
		record_error(result, "Unable to load file: %", path);
		return .{};
	}
	file_string, ok := read_entire_file(path);

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

	return make_file_from_string(file_string, path);
}

make_file_from_string :: (source : string, path : string = "") -> Input_File {
	input_file : Input_File;

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

	return input_file;
}

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);

			newline_after := type.children.count / 4;

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

				if it_index % newline_after == 0 {
					append(builder, "\n");
					indent(builder, 4);
				}
			}

			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 :: (ctx : *Compiler_Context, 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(ctx.scope_stack, variable.typename, xx 1);
			assert(find_result != null, "Internal compiler error\n");

			type_var := from_handle(ctx.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(ctx, child);
				array_add(*type.children, child_field);
			}

			type.name = variable.typename;
		}
	}

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

		if lookup_hint(hint.ident_value) == .Position {
			field_hint.kind = .Position;
		} else if lookup_hint(hint.ident_value) == .UV {
			field_hint.kind = .UV;
		} else if lookup_hint(hint.ident_value) == .Target  {
			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 {
			field_hint.custom_hint_name = hint.ident_value;
			field_hint.kind = .Custom;
		}
		array_add(*field.hints, field_hint);
	}

	field.type = type;

	return field;
}

type_variable_to_field :: (ctx : *Compiler_Context, variable : Type_Variable_Handle) -> Field {
	return type_variable_to_field(ctx, from_handle(ctx.type_variables, variable));
}

generate_output_data :: (ctx : *Compiler_Context) {
	if ctx.had_error {
		return;
	}
	
	if ctx.vertex_entry_point.node {
		ctx.vertex_entry_point.name = ctx.vertex_entry_point.node.name;

		type_variable := from_handle(ctx.type_variables, ctx.vertex_entry_point.node.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(ctx.type_variables, child.type_variable);
					field := type_variable_to_field(ctx, tv);
					array_add(*ctx.vertex_entry_point.input, field);
				}
			}
		}
	}

	for buffer_variable : ctx.constant_buffers {
		variable := from_handle(ctx.type_variables, buffer_variable);

		cb := array_add(*ctx.cbuffers);
		cb.name = variable.name;

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

		cb.buffer_index = variable.resource_index;

		for hint : variable.source_node.hint_tokens {
			field_hint : Field_Hint;
			field_hint.custom_hint_name = hint.ident_value;
			field_hint.kind = .Custom;
			array_add(*cb.hints, field_hint);	
		}
	}

	if ctx.pixel_entry_point.node {
		ctx.pixel_entry_point.name = ctx.pixel_entry_point.node.name;

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

		if type_variable.return_type_variable > 0 {
			field := type_variable_to_field(ctx, type_variable.return_type_variable);
			for hint : type_variable.source_node.hint_tokens {
				field_hint : Field_Hint;

				if lookup_hint(hint.ident_value) == .Position {
					field_hint.kind = .Position;
				} else if lookup_hint(hint.ident_value) == .Target {
					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);
			}
			ctx.pixel_entry_point.return_value = field;
		}		
	}
}

compile_file :: (ctx : *Compiler_Context, path : string, allocator : Allocator = temp) {
	new_context := context;
	new_context.allocator = allocator;
	push_context new_context {
		init_context_allocators();
		defer clear_context_allocators();

		ctx.file = make_file(ctx, path);

		lex(ctx, allocator);
		parse(ctx, allocator);
		check(ctx, allocator);
		codegen(ctx, allocator);
		generate_output_data(ctx);
	}
}