/////////////////////////////////////
//~ nbr: General improvements
//
// [x] Print out all failed tests in a list at the end
// [ ] Use unix (posix? bash? ascii?) color codes for errors
// [ ] Print golden file as green and new output as red

#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 :: "shd";
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, allocator := context.allocator) -> string {
	path := parse_path(file_path);
	file_without_extension := split(path.words[path.words.count - 1], ".");

	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;

	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);
	append(*builder, file_without_extension[0]);
	append(*builder, ".");
	append(*builder, GOLDEN_EXTENSION);
	golden_path := builder_to_string(*builder);
	array_add(*path.words, golden_path);

	final_path := path_to_string(path);

	return final_path;
}

run_lexer_test :: (file_path : string, lexer : *Lexer, output_type : Output_Type = 0) -> Result {
	ok := read_input_from_file(lexer, file_path);
	
	result_data : Result;
	result_data.path = file_path;
	result_data.stage = .Lexer;

	if !ok {
		result_data.type = .File_Read_Failed;
		result_data.info_text = tprint("Unable to read file: %\n", file_path);
		
		return result_data;
	} else {
		result_text : string;
		result := lex(lexer, *temp);

		if result.had_error {
			result_data.type = .Failed;
			result_text = report_messages(result.messages);
		} else {
			result_text = pretty_print_tokens(result.tokens, *temp);
		}

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

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

run_parser_test :: (file_path : string, output_type : Output_Type = 0) -> Result, *AST_Node {
	lexer : Lexer;
	result_data : Result;
	result_data.path = file_path;
	
	ok := read_input_from_file(*lexer, file_path);
	if !ok {
		log_error("Unable to read file: %\n", file_path);
		result_data.type = .File_Read_Failed;
		result_data.stage = .Lexer;
		return result_data, null;
	}

	result := lex(*lexer, *temp);
	if result.had_error {
		result_data.type = .Passed;
		return result_data, null;
	}

	result_data =, root := run_parser_test(*lexer, output_type);

	return result_data, root;
}

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

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

run_parser_test :: (lexer : *Lexer, output_type : Output_Type = 0) -> Result, *AST_Node {
	parse_state : Parse_State;
	result_data : Result;
	result_data.path = lexer.path;
	result_data.stage = .Parser;
	init_parse_state(*parse_state, lexer.result.tokens, lexer.path);

	result := parse(*parse_state);
	result_node : *AST_Node;
	result_text : string;
	
	if result.had_error {
		result_data.type = .Failed;
		result_text = report_messages(result.messages,, temp);
	} else {
		result_text = pretty_print_ast(parse_state.result.root, *temp);
		result_node = parse_state.result.root;
	}

	if output_type & .StdOut {
		result_data.info_text = result_text;
		result_data.type = .StdOut;
		return result_data, result_node;
	}
	
	golden_path := get_golden_path(parse_state.path, .Parser);
	do_golden_comparison(golden_path, result_text, *result_data, output_type);
	return result_data, result_node;
}

run_semantic_analysis_test :: (file_path : string, output_type : Output_Type = 0) -> Result, Semantic_Check_Result {
	lexer : Lexer;
	result_data : Result;
	result_data.path = file_path;
	
	ok := read_input_from_file(*lexer, file_path);
	if !ok {
		log_error("Unable to read file: %\n", file_path);
		result_data.type = .File_Read_Failed;
		result_data.stage = .Lexer;
		
		return result_data, .{};
	}

	lex_result := lex(*lexer, *temp);
	if lex_result.had_error {
		result_data.type = .Failed;
		result_data.stage = .Lexer;
		result_data.info_text = report_messages(lex_result.messages);

		if output_type & .StdOut {
			result_data.type = .StdOut;
			return result_data, .{};
		}
		golden_path := get_golden_path(file_path, .Semantic_Analysis);
		do_golden_comparison(golden_path, result_data.info_text, *result_data, output_type);
		return result_data, .{};
	}

	parse_state : Parse_State;
	result_data.stage = .Parser;
	init_parse_state(*parse_state, lex_result.tokens, lexer.path);

	parse_result := parse(*parse_state);
	if parse_result.had_error {
		result_data.type = .Failed;
		result_data.info_text = report_messages(parse_result.messages);

		if output_type & .StdOut {
			result_data.type = .StdOut;
			return result_data, .{};
		}
		
		golden_path := get_golden_path(file_path, .Semantic_Analysis);
		do_golden_comparison(golden_path, result_data.info_text, *result_data, output_type);
		return result_data, .{};
	}

	result, check_result := run_semantic_analysis_test(file_path, parse_state.result.root, output_type);
	return result, check_result;
}

run_semantic_analysis_test :: (file_path : string, root : *AST_Node, output_type : Output_Type = 0) -> Result, Semantic_Check_Result {
	result_data : Result;

	result_data.path = file_path;
	result_data.stage = .Semantic_Analysis;
	checker : Semantic_Checker;
	init_semantic_checker(*checker, root, file_path);
	result_text : string;

	result := check(*checker);
	if result.had_error {
		result_data.type = .Failed;
		result_text = report_messages(checker.result.messages);
	} else {
		result_text = pretty_print_symbol_table(*checker, temp);
		constraints := pretty_print_type_constraints(*checker, temp);
		type_vars := pretty_print_type_variables(*checker, temp);
		// print("Constraints\n%\n", constraints);
		// print("Solution\n%\n", type_vars);		
	}

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

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

run_codegen_test :: (path : string, root : *AST_Node, check_result : Semantic_Check_Result, output_type : Output_Type = 0) -> Result, Codegen_Result {

	result_data : Result;
	result_data.path = path;
	result_data.stage = .Codegen;

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

	result_text : string;

	result := codegen(*state);
	
	if result.had_error {
		result_data.type = .Failed;
		result_data.info_text = report_messages(result.messages);
		return result_data, .{};
	}
	
	result_text = result.result_text;

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

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

run_codegen_test :: (path : string, root : *AST_Node, output_type : Output_Type = 0) -> Result, Codegen_Result {
	checker : Semantic_Checker;
	init_semantic_checker(*checker, root, path);

	result_data : Result;
	result_data.path = path;
	result_data.stage = .Semantic_Analysis;
	
	check_result := check(*checker);
	if check_result.had_error {
		result_data.type = .Failed;
		result_data.info_text = report_messages(check_result.messages);
		return result_data, .{};
	}

	result, codegen_result := run_codegen_test(path, root, check_result, output_type);
	return result, codegen_result;
}

run_codegen_test :: (path : string, output_type : Output_Type = 0) -> Result, Codegen_Result {
	lexer : Lexer;
	result_data : Result;
	result_data.path = path;

	ok := read_input_from_file(*lexer, path);
	if !ok {
		log_error("Unable to read file: %\n", path);
		result_data.type = .File_Read_Failed;
		result_data.stage = .Lexer;
		return result_data, .{};
	}

	lex_result := lex(*lexer, *temp);
	if lex_result.had_error {
		result_data.type = .Failed;
		result_data.stage = .Lexer;
		return result_data, .{};
	}

	parse_state : Parse_State;
	result_data.stage = .Parser;
	init_parse_state(*parse_state, lex_result.tokens, lexer.path);

	parse_result := parse(*parse_state);
	if parse_result.had_error {
		result_data.type = .Failed;
		result_data.info_text = pretty_print_ast(parse_result.root, *temp);
		return result_data, .{};
	}

	result, codegen_result := run_codegen_test(path, parse_result.root, output_type);
	return result, codegen_result;
}

run_compile_test :: (path : string, output_type : Output_Type = 0) -> Result, Compilation_Result {
	compiler : Shader_Compiler;
	result : Result;
	compilation_result := compile_file(*compiler, path);
	if compilation_result.had_error {
		result.type = .Failed;
		result.info_text = tprint("Failed compiling: %\n", path);
	}

	return result, compilation_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 :: (file_path : string, stage_flags : Stage_Flags, results : *[..]Result, output_type : Output_Type = 0) {
	lexer : Lexer;
	result : Result;
	if stage_flags & .Lexer {
		result = run_lexer_test(file_path, *lexer, output_type);
		record_result(results, result);
	}

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

	check_result : Semantic_Check_Result;
	if stage_flags & .Semantic_Analysis {
		if stage_flags & .Parser && (result.type == .Passed || result.type == .Golden_Output) {
			result, check_result = run_semantic_analysis_test(file_path, root_node, output_type);
		} else {
			result, check_result = run_semantic_analysis_test(file_path, 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(file_path, root_node, check_result, output_type);
		} else if root_node {
			result = run_codegen_test(file_path, root_node, output_type);
		} else {
			result = run_codegen_test(file_path, output_type);
		}
		record_result(results, result);
	}

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

run_test :: (test_case : Test_Case, results : *[..]Result, output_type : Output_Type = 0) {
	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(test_case.path, test_case.stage_flags, results, output_type);
}

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;
	}
	
	failed_test_paths : [..]Fail_Data;
	failed_test_paths.allocator = temp;

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

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

		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");
	}
	print("\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));
}

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

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

	for split_line : split_lines {
		line := split(split_line, " ");
		if line[0].count == 0 {
			continue;
		}
		
		if line[0].data[0] == #char "#" {
			continue;
		}
		
		if line.count == 1 {
			log_error("Invalid line - % - %\n", it_index + 1, line);
			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);
		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 :: () {
	lexer : Lexer;

	args := get_command_line_arguments();

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

	Argument_Parse_State :: enum {
		None;
		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);
				}
			}
			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, ".") {
					split_path := split(arg, ".");
					extension := split_path[1];
					if extension == SHADER_EXTENSION {
						path := copy_string(arg);
						test_case := make_test_case(path, 0);
						array_add(*current_suite.test_cases, test_case);
					}
				} else {
					print("%Unknown argument %\n", red, arg);
				}
			}
			case .None; {
				if contains(arg, ".") {
					split_path := split(arg, ".");
					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;
							array_add(*suites, suite);
							current_suite = *suites[0];
						}
						
						arg_parse_state = .Run_Test;
						path := copy_string(arg);
						test_case := make_test_case(path, 0);
						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;
						read_suite(path, *suite);
						array_add(*suites, suite);
						current_suite = *suites[0];
					}
				}			
			}
		}
	}

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