/////////////////////////////////////
/*~ nbr: General improvements
- [x] Print out all failed tests in a list at the end
- [x] Use new compiler API with Compile_Result and Compiled_File instead
- [ ] Use unix (posix? bash? ascii?) color codes for errors
- [ ] Print golden file as green and new output as red
- [ ] Rename to Ink.jai
- [ ] Add -test option. -test does the same as test.exe used to do
- [ ] Add -fuzz option to run fuzzer (add args later)
- [ ] Add -output option to output the compiled file. Issue with this is the generated data can't be output like that. Would require serialization.
*/

#import "Basic";
#import "File";
#import "String";
#import "File_Utilities";
#import "Print_Color";

#load "module.jai";

GOLDEN_EXTENSION         :: "golden";
LEXER_FOLDER             :: "lex";
PARSER_FOLDER            :: "parse";
CODEGEN_FOLDER           :: "codegen";
COMPILED_FOLDER          :: "compiled";
SEMANTIC_ANALYSIS_FOLDER :: "semant";
TESTS_FOLDER             :: "test";

SHADER_EXTENSION :: "ink";
SUITE_EXTENSION  :: "suite";

Stage_Flags :: enum_flags u16 {
	Lexer             :: 0x1;
	Parser            :: 0x2;
	Semantic_Analysis :: 0x4;
	Codegen           :: 0x8;
	Compile           :: 0x10;
}

Output_Type :: enum_flags u16 {
	Golden :: 0x1;
	StdOut :: 0x2;
}

Result_Type :: enum {
	File_Read_Failed;
	Golden_File_Not_Found;
	StdOut;
	Golden_Output;
	Passed;
	Failed;
}

Result :: struct {
	type : Result_Type;
	path : string;
	stage : Stage_Flags;

	golden_path : string;
	info_text   : string;
}

Test_Case :: struct {
	path : string;
	stage_flags : Stage_Flags;
}

Test_Suite :: struct {
	name : string;
	test_cases : [..]Test_Case;

	results : [..]Result;
}

get_golden_path :: (file_path : string, stage : Stage_Flags) -> string {
	sc := get_scratch();
	defer scratch_end(sc);
	path := parse_path(file_path,, sc.allocator);
	file_without_extension := split(path.words[path.words.count - 1], ".",, sc.allocator);

	builder : String_Builder;
	builder.allocator = temp;

	final_path_length := file_path.count - SHADER_EXTENSION.count + GOLDEN_EXTENSION.count + 1;  // +1 for dot

	path.words.count -= 1;
	path.words.allocator = sc.allocator;

	if stage == {
		case .Lexer; {
			dir := tprint("%/%", TESTS_FOLDER, LEXER_FOLDER);
			make_directory_if_it_does_not_exist(dir);
			array_add(*path.words, LEXER_FOLDER);
		}
		case .Parser; {
			dir := tprint("%/%", TESTS_FOLDER, PARSER_FOLDER);
			make_directory_if_it_does_not_exist(dir);
			array_add(*path.words, PARSER_FOLDER);
		}
		case .Semantic_Analysis; {
			dir := tprint("%/%", TESTS_FOLDER, SEMANTIC_ANALYSIS_FOLDER);
			make_directory_if_it_does_not_exist(dir);
			array_add(*path.words, SEMANTIC_ANALYSIS_FOLDER);
		}
		case .Codegen; {
			dir := tprint("%/%", TESTS_FOLDER, CODEGEN_FOLDER);
			make_directory_if_it_does_not_exist(dir);
			array_add(*path.words, CODEGEN_FOLDER);
		}
		case .Compile; {
			dir := tprint("%/%", TESTS_FOLDER, COMPILED_FOLDER);
			make_directory_if_it_does_not_exist(dir);
			array_add(*path.words, COMPILED_FOLDER);
		}
	}

	init_string_builder(*builder, file_without_extension.count + GOLDEN_EXTENSION.count + 1);
	builder.allocator = sc.allocator;
	append(*builder, file_without_extension[0]);
	append(*builder, ".");
	append(*builder, GOLDEN_EXTENSION);
	golden_path := builder_to_string(*builder,, sc.allocator);
	array_add(*path.words, golden_path);

	final_path := path_to_string(path);

	return final_path;
}

do_golden_comparison :: (golden_path : string, comparison_text : string, result : *Result, output_type : Output_Type) {
	sc := get_scratch();
	defer scratch_end(sc);
	if output_type & .Golden {
		// Output the comparison file
		write_entire_file(golden_path, comparison_text);
		result.golden_path = copy_string(golden_path);
		result.type = .Golden_Output;
		return;
	} else {
		// Do the comparison
		if !file_exists(golden_path) {
			result.info_text = tprint("Golden file % does not exist. Please run with -output-as-golden at least once.\n", golden_path);
			result.type = .Golden_File_Not_Found;
			return;
		}
		
		golden_text, ok := read_entire_file(golden_path,, sc.allocator);
		if !ok {
			result.info_text = tprint("Unable to open golden file %\n", golden_path);
			result.type = .Golden_File_Not_Found;
			return;
		}

		comp := replace(comparison_text, "\r\n", "\n",, sc.allocator);
		gold := replace(golden_text, "\r\n", "\n",, sc.allocator);
		ok = compare(comp, gold) == 0;
		if !ok {
			result.type = .Failed;
			result.info_text = tprint("Golden file:\n%\n===============\n%", gold, comp);
		} else {
			result.type = .Passed;
		}
	}
}

run_codegen_test :: (file_path : string, ctx : *Compiler_Context, output_type : Output_Type = 0) -> Result {
	result : Result;
	result.path = file_path;

	lex(ctx, context.allocator);
	parse(ctx, context.allocator);
	check(ctx, context.allocator);

	if ctx.had_error {
		result.type = .Failed;
		return result;
	}
	result = run_codegen_test(ctx, output_type);

	return result;
}

run_codegen_test :: (ctx : *Compiler_Context, output_type : Output_Type = 0) -> Result {
	result : Result;
	result.path = ctx.file.path;
	result_text : string;
	
	codegen(ctx, context.allocator);

	if ctx.had_error {
		result.type = .Failed;
		result_text = report_messages(ctx.messages);
		return result;
	} 

	result_text = ctx.codegen_result_text;

	if output_type & .StdOut {
		result.info_text = result_text;
		result.type = .StdOut;
		return result;
	}

	golden_path := get_golden_path(ctx.file.path, .Codegen);
	do_golden_comparison(golden_path, result_text, *result, output_type);
	return result;
}

run_compile_test :: (path : string, output_type : Output_Type = 0) -> Result, Compiler_Context {
	ctx : Compiler_Context;
	result : Result;
	result.path = path;
	compile_file(*ctx, path, context.allocator);
	
	if ctx.had_error {
		result.type = .Failed;
		result.info_text = tprint("Failed compiling: %\n", path);
	} else {
		sc := get_scratch();
		defer scratch_end(sc);
		sb : String_Builder;
		init_string_builder(*sb,, sc.allocator);
		if ctx.vertex_entry_point.name.count > 0 {
			print_to_builder(*sb, "[vertex entry point] - %\n", ctx.vertex_entry_point.name);
		}
		if ctx.pixel_entry_point.name.count > 0 {
			print_to_builder(*sb, "[pixel entry point] - %\n", ctx.pixel_entry_point.name);
		}

		for cb : ctx.cbuffers {
			print_to_builder(*sb, "[constant_buffer] - % - %\n", cb.name, cb.buffer_index);

			indent(*sb, 1);
			for field : cb.fields {
				append(*sb, "[field] - ");
				pretty_print_field(*sb, *field.base_field);
			}
		}

		result.info_text = builder_to_string(*sb);
	}

	if output_type & .StdOut {
		result.type = .StdOut;
		return result, ctx;
	}

	golden_path := get_golden_path(ctx.file.path, .Compile);
	do_golden_comparison(golden_path, result.info_text, *result, output_type);

	return result, ctx;
}

run_lexer_test :: (file_path : string, ctx : *Compiler_Context, output_type : Output_Type = 0) -> Result {
	result : Result;
	result.path = file_path;
	result.stage = .Lexer;

	result_text : string;
	
	lex(ctx);
	if ctx.had_error {
		result.type = .Failed;
		result_text = report_messages(ctx.messages);
	} else {
		result_text = pretty_print_tokens(ctx.tokens, context.allocator);
	}

	if output_type & .StdOut {
		result.info_text = result_text;
		result.type = .StdOut;
		return result;
	}

	golden_path := get_golden_path(file_path, .Lexer);
	do_golden_comparison(golden_path, result_text, *result, output_type);
	return result;
}

run_parser_test :: (file_path : string, ctx : *Compiler_Context, output_type : Output_Type = 0) -> Result {
	result : Result;
	result.path = file_path;

	lex(ctx);
	if ctx.had_error {
		result.type = .Passed;
		return result;
	}

	result = run_parser_test(ctx, output_type);

	return result;
}

run_parser_test :: (ctx : *Compiler_Context, output_type : Output_Type = 0) -> Result {
	parse(ctx, context.allocator);
	result : Result;
	result.path = ctx.file.path;
	result_text : string;

	if ctx.had_error {
		result.type = .Failed;
		result_text = report_messages(ctx.messages);
	} else {
		result_text = pretty_print_ast(ctx.root, context.allocator);
	}

	if output_type & .StdOut {
		result.info_text = result_text;
		result.type = .StdOut;
		return result;
	}
	
	golden_path := get_golden_path(ctx.file.path, .Parser);
	do_golden_comparison(golden_path, result_text, *result, output_type);
	return result;
}

run_semantic_analysis_test :: (ctx : *Compiler_Context, output_type : Output_Type = 0) -> Result {
	result : Result;
	result.path = ctx.file.path;
	result_text : string;
	
	check(ctx, context.allocator);

	if ctx.had_error {
		result.type = .Failed;
		result_text = report_messages(ctx.messages);
	} else {
		result_text = pretty_print_symbol_table(ctx, context.allocator);
	}

	if output_type & .StdOut {
		result.info_text = result_text;
		result.type = .StdOut;
		return result;
	}

	golden_path := get_golden_path(ctx.file.path, .Semantic_Analysis);
	do_golden_comparison(golden_path, result_text, *result, output_type);
	return result;
}

run_semantic_analysis_test :: (file_path : string, ctx : *Compiler_Context, output_type : Output_Type = 0) -> Result {
	result : Result;
	result.path = file_path;

	lex(ctx, context.allocator);
	parse(ctx, context.allocator);
	if ctx.had_error {
		result.type = .Failed;
		return result;
	}
	
	result = run_semantic_analysis_test(ctx, output_type);

	return result;
}

make_test_case :: (path : string, stage_flags : Stage_Flags, allocator := context.allocator) -> Test_Case {
	test_case : Test_Case;
	test_case.path = copy_string(path,, allocator);
	replace_chars(test_case.path, "\\", #char "/");
	test_case.stage_flags = stage_flags;

	return test_case;
}

run_test_new :: (file_path : string, stage_flags : Stage_Flags, results : *[..]Result, output_type : Output_Type = 0, allocator := temp) {
	new_context := context;
	new_context.allocator = allocator;
	push_context new_context {
		ctx : Compiler_Context;

		ctx.file = make_file(*ctx, file_path);
		
		result : Result;
		if stage_flags & .Lexer {
			result = run_lexer_test(file_path, *ctx, output_type);
			record_result(results, result);
		}

		if stage_flags & .Parser {
			if stage_flags & .Lexer && result.type == .Passed || result.type == .Golden_Output {
				result = run_parser_test(*ctx, output_type);
			} else {
				result = run_parser_test(file_path, *ctx, output_type);
			}
			record_result(results, result);
		}

		if stage_flags & .Semantic_Analysis {
			if stage_flags & .Parser && (result.type == .Passed || result.type == .Golden_Output) {
				result = run_semantic_analysis_test(*ctx, output_type);
			} else {
				result = run_semantic_analysis_test(file_path, *ctx, output_type);
			}
			record_result(results, result);
		}

		if stage_flags & .Codegen {
			if stage_flags & .Semantic_Analysis && (result.type == .Passed || result.type == .Golden_Output) {
				result = run_codegen_test(*ctx, output_type);
			} else {
				result = run_codegen_test(file_path, *ctx, output_type);
			}
			record_result(results, result);
		}

		if stage_flags & .Compile {
			result = run_compile_test(file_path, output_type);
			record_result(results, result);
		}
	}
	
}

run_test :: (test_case : Test_Case, results : *[..]Result, output_type : Output_Type = 0, allocator := temp) {
	print("%Running test: %......", cyan(), test_case.path);

	// path 30
	// len 35
	// == 5


	// path 20
	// len = 35
	// == 15
	
	len := 50;
	rest := len - test_case.path.count;
	for i: 0..rest {
		print(" ");
	}

	run_test_new(test_case.path, test_case.stage_flags, results, output_type, allocator);
}

record_result :: (results : *[..]Result, result : Result) {
	array_add(results, result);
}

run_test_suite :: (using suite : *Test_Suite, output_type : Output_Type = 0) {
	if suite.name.count > 0 {
		print("%Running suite: %\n", green(), suite.name);
		print("%", reset_color());
	}

	Fail_Data :: struct {
		path : string;
		stage : string;
	}

	test_arena : Allocator = make_arena(Gigabytes(1));
	
	failed_test_paths : [..]Fail_Data;
	failed_test_paths.allocator = test_arena;

	builder : String_Builder;
	init_string_builder(*builder,, test_arena);

	for test_case : test_cases {
		run_test(test_case, *suite.results, output_type, allocator = test_arena);

		for < suite.results {
			result := suite.results[it_index];
			if compare(result.path, test_case.path) == 0 {
				if result.type == {
					case .Failed; {
						array_add(*failed_test_paths, .{ result.path, stage_to_string(result.stage) });
					}
					case .File_Read_Failed; {
						array_add(*failed_test_paths, .{ result.path, "file not found" });
					}
					case .Golden_File_Not_Found; {
						array_add(*failed_test_paths, .{ result.path, tprint("golden file not found for %", stage_to_string(result.stage)) });
					}
				}
				evaluate_result(result);
			} else {
				break;
			}
		}

		// print("\n");
	}
	append(*builder, "\n");

	if output_type == 0 {
		if failed_test_paths.count == 0 {
			green(*builder);
			print_to_builder(*builder, "All % tests passed!\n", test_cases.count);
			reset_color(*builder);
		} else {
			print_to_builder(*builder, "%/% tests passed\n", test_cases.count - failed_test_paths.count, test_cases.count);
			red(*builder);

			print_to_builder(*builder, "% failed\n", failed_test_paths.count);
			for failed_test : failed_test_paths {
				print_to_builder(*builder, "% failed with error: %\n", failed_test.path, failed_test.stage);
			}
			reset_color(*builder);
		}
	} 

	print("%\n", builder_to_string(*builder,, test_arena));
}

read_suite :: (file_path : string, suite : *Test_Suite, allocator := temp) -> bool {
	sc := get_scratch();
	defer scratch_end(sc);
	bytes, ok := read_entire_file(file_path,, sc.allocator);
	if !ok {
		log_error("Unable to read suite file %\n", file_path);
		return false;
	}

	path := parse_path(file_path,, sc.allocator);
	file_without_extension := split(path.words[path.words.count - 1], ".",, sc.allocator);
	suite.name = copy_string(file_without_extension[0],, allocator);
	split_lines := split(bytes, "\n",, sc.allocator);

	for split_line : split_lines {
		if split_line.count == 0 {
			break;
		}
		
		if split_line[0] == #char "#" {
			continue;
		}
		
		line := split(split_line, " ",, sc.allocator);
		if line[0].count == 0 {
			continue;
		}
		
		if line[0].data[0] == #char "#" {
			continue;
		}
		
		if line.count == 1 {
			line = split(split_line, "\t",, sc.allocator);
			if line.count == 1 {
				log_error("Invalid line - % - \n", it_index + 1);
				continue;
			}
		}
		test_case_path := line[0];
		stage_flags : Stage_Flags;

		for i: 0..line.count - 1 {
			trimmed := trim(line[i]);
			if equal(trimmed, "lex") {
				stage_flags |= .Lexer;
			} else if equal(trimmed, "parse") {
				stage_flags |= .Parser;
			} else if equal(trimmed, "semant") {
				stage_flags |= .Semantic_Analysis;
			} else if equal(trimmed, "codegen") {
				stage_flags |= .Codegen;
			} else if equal(trimmed, "compile") {
				stage_flags |= .Compile;
			}
		}
		test_case := make_test_case(test_case_path, stage_flags, allocator);
		array_add(*suite.test_cases, test_case);
	}

	return true;
}

read_test :: () {

}

stage_to_string :: (stage : Stage_Flags) -> string {
	if #complete stage == {
		case .Lexer; return "lexing";
		case .Parser; return "parsing";
		case .Semantic_Analysis; return "semantic checking";
		case .Codegen; return "codegen";
		case .Compile; return "compiled";
		case; return "";
	}
}

evaluate_result :: (result : Result) {
	stage : string = stage_to_string(result.stage);
	
	if #complete result.type == {
		case .File_Read_Failed; {
			print(" %", red());
			print("failed with File_Read_Failed\n");
		}
		case .Golden_File_Not_Found; {
			print(" %", red());
			print("failed with Golden File Not Found for stage %\n", stage);
		}
		case .StdOut; {
		}
		case .Golden_Output; {
			print(" %", yellow());
			print("output new golden file at %\n", result.golden_path);
		}
		case .Passed; {
			print(" %", green());
			print("passed %\n", stage);
		}
		case .Failed; {
			print(" %", red());
			print("failed %\n", stage);
		}
	}

	if result.info_text.count > 0 {
		print("%", cyan());
		print("--- Info text ---\n");
		print("%", yellow());
		print("%\n", result.info_text);
	}

	print("%", reset_color());
}

main :: () {
	args := get_command_line_arguments();

	init_context_allocators();

	local_temp := make_arena(Megabytes(128));

	suites : [..]Test_Suite;
	suites.allocator = local_temp;
	output_type : Output_Type = 0;

	Argument_Parse_State :: enum {
		None;
		Compile;
		Run_Suite;
		Run_Test;
	}

	arg_parse_state : Argument_Parse_State;
	current_suite : *Test_Suite;

	path : string;

	for i: 1..args.count - 1 {
		arg := args[i];
		if arg == "-output-as-golden" {
			output_type |= .Golden;
			continue;
		} else if arg == "-output" {
			output_type |= .StdOut;
			continue;
		}

		if arg_parse_state == {
			case .Run_Suite; {
				if arg == "-output-as-golden" {
					output_type |= .Golden;
				} else if arg == "-output" {
					output_type |= .StdOut;
				} else {
					print("%Unknown argument % %\n", red(), arg, reset_color());
				}
			}
			case .Run_Test; {
				cases := current_suite.test_cases.count;
				if arg == "-lex" {
					current_suite.test_cases[cases - 1].stage_flags |= .Lexer;
				} else if arg == "-parse" {
					current_suite.test_cases[cases - 1].stage_flags |= .Parser;
				} else if arg == "-semant" {
					current_suite.test_cases[cases - 1].stage_flags |= .Semantic_Analysis;
				} else if arg == "-codegen" {
					current_suite.test_cases[cases - 1].stage_flags |= .Codegen;
				} else if arg == "-compile" {
					current_suite.test_cases[cases - 1].stage_flags |= .Compile;
				} else if contains(arg, ".") {
					sc := get_scratch();
					defer scratch_end(sc);
					path_split := split(arg, "\\",, sc.allocator);
					split_path := split(path_split[path_split.count - 1], ".",, sc.allocator);
					extension := split_path[1];
					if extension == SHADER_EXTENSION {
						path := copy_string(arg,, local_temp);
						test_case := make_test_case(path, 0, local_temp);
						array_add(*current_suite.test_cases, test_case);
					} else {
						print("%Invalid file as argument % %\n", red(), arg, reset_color());						
					}
				} else {
					print("%Unknown argument % %\n", red(), arg, reset_color());
				}
			}
			case .None; {
				if contains(arg, ".") {
					sc := get_scratch();
					defer scratch_end(sc);
					path_split := split(arg, "\\",, sc.allocator);
					split_path := split(path_split[path_split.count - 1], ".",, sc.allocator);
					extension := split_path[1];

					if extension == SHADER_EXTENSION {
						if arg_parse_state == .Run_Suite {
							log_error("Unable to run a test while already running suite.");
							continue;
						}

						if !current_suite {
							suite : Test_Suite;
							suite.results.allocator = local_temp;
							suite.test_cases.allocator = local_temp;
							array_add(*suites, suite);
							current_suite = *suites[0];
						}
						arg_parse_state = .Run_Test;
						path := copy_string(arg,, local_temp);
						test_case := make_test_case(path, 0, local_temp);
						array_add(*current_suite.test_cases, test_case);
					} else if extension == SUITE_EXTENSION {
						if arg_parse_state == .Run_Test {
							log_error("Unable to run a suite while already running test.");
							continue;
						}
						arg_parse_state = .Run_Suite;
						path := copy_string(arg);

						suite : Test_Suite;
						suite.results.allocator = local_temp;
						suite.test_cases.allocator = local_temp;
						read_suite(path, *suite, local_temp);
						array_add(*suites, suite);
						current_suite = *suites[0];
					} else {
						print("%Invalid file as argument % %\n", red(), arg, reset_color());
					}
				}
			}
		}
	}

	for suite : suites {
		run_test_suite(*suite, output_type);
	}

	clear(local_temp);
}