package compiler

import "./ast"
import "./lexer"
import "./parser"
import "./source"
import "./symbol"
import "./types"
import "core:fmt"
import "core:path/filepath"
import "core:slice"
import "core:strings"

Test_Entry :: struct {
	function:     ast.Function_Id,
	package_path: string,
	source_path:  string,
	offset:       source.Offset,
}

test_entry_less :: proc(a, b: Test_Entry) -> bool {
	if a.package_path != b.package_path {
		return a.package_path < b.package_path
	}
	if a.source_path != b.source_path {
		return a.source_path < b.source_path
	}
	return a.offset < b.offset
}

testing_package :: proc(module: ^ast.Module, project_root: string) -> ast.Package_Id {
	path, err := filepath.join({project_root, "std", "testing"}, module.allocator)
	if err != nil {
		return ast.INVALID_PACKAGE
	}
	defer delete(path, module.allocator)
	canonical, ok := filepath.abs(path, module.allocator)
	if !ok {
		return ast.INVALID_PACKAGE
	}
	defer delete(canonical, module.allocator)
	for pkg, index in module.packages {
		if pkg.path == canonical {
			return ast.package_id(index)
		}
	}
	return ast.INVALID_PACKAGE
}

source_file_id :: proc(module: ^ast.Module, id: source.Source_Id) -> ast.File_Id {
	for file, index in module.files {
		if file.source == id {
			return ast.file_id(index)
		}
	}
	return ast.INVALID_FILE
}

append_ast_expr :: proc(module: ^ast.Module, expr: ast.Expr) -> ast.Expr_Id {
	id := ast.expr_id(len(module.exprs))
	append(&module.exprs, expr)
	return id
}

location_expr :: proc(
	module: ^ast.Module,
	sources: ^source.Store,
	symbols: ^symbol.Table,
	span: source.Span,
	qualifier: symbol.Id,
) -> ast.Expr_Id {
	line, column := 1, 1
	path := "<unknown>"
	if int(span.file) < len(sources.items) {
		file := &sources.items[span.file]
		line, column = source.line_and_column(file, span.start)
		path = file.path
	}
	string_index := u64(len(module.strings))
	append(&module.strings, strings.clone(path, module.allocator))
	file_expr := append_ast_expr(module, ast.Expr{
		kind=.String, span=span, integer=string_index,
		left=ast.INVALID_EXPR, right=ast.INVALID_EXPR,
		diagnostic=source.INVALID_DIAGNOSTIC,
	})
	line_expr := append_ast_expr(module, ast.Expr{
		kind=.Integer, span=span, integer=u64(line),
		left=ast.INVALID_EXPR, right=ast.INVALID_EXPR,
		diagnostic=source.INVALID_DIAGNOSTIC,
	})
	column_expr := append_ast_expr(module, ast.Expr{
		kind=.Integer, span=span, integer=u64(column),
		left=ast.INVALID_EXPR, right=ast.INVALID_EXPR,
		diagnostic=source.INVALID_DIAGNOSTIC,
	})
	values := []ast.Expr_Id{file_expr, line_expr, column_expr}
	names := []string{"file", "line", "column"}
	fields := make([]ast.Expr_Id, 3, module.allocator)
	for index in 0..<3 {
		fields[index] = append_ast_expr(module, ast.Expr{
			kind=.Keyed, span=span,
			name=symbol.intern(symbols, names[index]),
			left=values[index], right=ast.INVALID_EXPR,
			diagnostic=source.INVALID_DIAGNOSTIC,
		})
	}
	return append_ast_expr(module, ast.Expr{
		kind=.Struct_Literal, span=span,
		qualifier=qualifier,
		name=symbol.intern(symbols, "SourceLocation"),
		args=fields,
		left=ast.INVALID_EXPR, right=ast.INVALID_EXPR,
		diagnostic=source.INVALID_DIAGNOSTIC,
	})
}

inject_assertion_locations :: proc(
	module: ^ast.Module,
	sources: ^source.Store,
	symbols: ^symbol.Table,
	testing_pkg: ast.Package_Id,
) {
	expect := symbol.intern(symbols, "expect")
	expect_equal := symbol.intern(symbols, "expect_equal")
	expect_type := symbol.intern(symbols, "expect_type")
	original_count := len(module.exprs)
	for index in 0..<original_count {
		expr := &module.exprs[index]
		if expr.kind != .Call || expr.intrinsic || !symbol.is_valid(expr.qualifier) ||
		   (expr.name != expect && expr.name != expect_equal && expr.name != expect_type) {
			continue
		}
		file := source_file_id(module, expr.span.file)
		matched := false
		for import_item in module.imports {
			if !import_item.test_only && import_item.file == file &&
			   import_item.alias == expr.qualifier && import_item.target == testing_pkg {
				matched = true
				break
			}
		}
		if !matched {
			continue
		}
		location := location_expr(module, sources, symbols, expr.span, expr.qualifier)
		expr = &module.exprs[index]
		args := make([]ast.Expr_Id, len(expr.args)+1, module.allocator)
		copy(args, expr.args)
		args[len(expr.args)] = location
		delete(expr.args, module.allocator)
		expr.args = args
	}
}

write_brolang_string :: proc(builder: ^strings.Builder, value: string) {
	for byte_value in transmute([]byte)value {
		if byte_value == '\\' || byte_value == '"' {
			strings.write_byte(builder, '\\')
		}
		strings.write_byte(builder, byte_value)
	}
}

append_runner :: proc(
	module: ^ast.Module,
	sources: ^source.Store,
	diagnostics: ^source.Diagnostics,
	symbols: ^symbol.Table,
	testing_pkg: ast.Package_Id,
	testing_error: types.Type,
	tests: []Test_Entry,
) {
	builder := strings.builder_make(module.allocator)
	defer strings.builder_destroy(&builder)
	for entry, index in tests {
		test := module.functions[entry.function]
		alias := fmt.tprintf("__brolang_test_%d", index)
		fmt.sbprintf(&builder, "@hide __brolang_test_adapter_%d func() void ! __brolang_testing.Error ", index)
		strings.write_string(&builder, "{\n\t")
		fmt.sbprintf(&builder, "%s.%s() catch |_| ", alias, symbol.resolve(symbols, test.name))
		strings.write_string(&builder, "{\n\t\treturn .expectation_failed\n\t}\n}\n\n")
	}
	strings.write_string(&builder, "main func() i32 {\n\tfailed i32 := 0\n")
	for entry, index in tests {
		test_id := entry.function
		test := module.functions[test_id]
		pkg := module.packages[test.pkg]
		name := fmt.tprintf("%s.%s", filepath.base(pkg.path), symbol.resolve(symbols, test.name))
		strings.write_string(&builder, "\tif (!__brolang_testing.run(\"")
		write_brolang_string(&builder, name)
		fmt.sbprintf(&builder, "\", __brolang_test_adapter_%d)) ", index)
		strings.write_string(&builder, "{\n\t\tfailed += 1\n\t}\n")
	}
	fmt.sbprintf(&builder, "\t__brolang_testing.summary(%d - failed, failed)\n", len(tests))
	strings.write_string(&builder, "\tif (failed != 0) return 1\n\treturn 0\n}\n")

	runner_text := strings.to_string(builder)
	source_id := source.add_source(sources, "<brolang-test-runner>", runner_text)
	file_id := ast.file_id(len(module.files))
	append(&module.files, ast.File{source=source_id, pkg=0})
	stream := lexer.lex(&sources.items[source_id], diagnostics, symbols, module.allocator)
	defer delete(stream.items)
	function_start := len(module.functions)
	parser.parse_into(&stream, &sources.items[source_id], diagnostics, module, 0, file_id)
	for index in 0..<len(tests) {
		module.functions[function_start+index].error = testing_error
	}

	append(&module.imports, ast.Import{
		alias=symbol.intern(symbols, "__brolang_testing"),
		path=strings.clone("@std/testing", module.allocator),
		pkg=0, file=file_id, target=testing_pkg,
		valid=true, used=true,
		diagnostic=source.INVALID_DIAGNOSTIC,
	})
	for entry, index in tests {
		test_id := entry.function
		test := module.functions[test_id]
		append(&module.imports, ast.Import{
			alias=symbol.intern(symbols, fmt.tprintf("__brolang_test_%d", index)),
			path=strings.clone(module.packages[test.pkg].path, module.allocator),
			pkg=0, file=file_id, target=test.pkg,
			valid=true, used=true,
			diagnostic=source.INVALID_DIAGNOSTIC,
		})
	}
}

prepare_tests :: proc(
	module: ^ast.Module,
	sources: ^source.Store,
	diagnostics: ^source.Diagnostics,
	symbols: ^symbol.Table,
	mode: ast.Compile_Mode,
	project_root: string,
) -> bool {
	if mode == .Executable {
		for &function in module.functions {
			if !function.test {
				continue
			}
			function.generated = true
			for &import_item in module.imports {
				if !import_item.test_only && import_item.file == function.file {
					import_item.used = true
				}
			}
		}
		return true
	}

	testing_pkg := testing_package(module, project_root)
	if testing_pkg == ast.INVALID_PACKAGE {
		source.add(diagnostics, source.Span{}, "test builds require @std/testing")
		return false
	}
	error_type := types.find_named(
		&module.type_store,
		u32(testing_pkg),
		u32(symbol.intern(symbols, "Error")),
	)
	if !types.is_enum(error_type, &module.type_store) {
		source.add(diagnostics, source.Span{}, "@std/testing must declare Error as an enum")
		return false
	}

	tests: [dynamic]Test_Entry
	tests.allocator = module.allocator
	defer delete(tests)
	main_name := symbol.intern(symbols, "main")
	for &function, index in module.functions {
		if function.test {
			if int(function.pkg) < len(module.packages) && module.packages[function.pkg].test {
				function.result = types.VOID
				function.error = error_type
				file := module.files[function.file]
				append(&tests, Test_Entry{
					function=ast.function_id(index),
					package_path=module.packages[function.pkg].path,
					source_path=sources.items[file.source].path,
					offset=function.span.start,
				})
			} else {
				function.generated = true
			}
		} else if function.pkg == 0 && function.name == main_name {
			function.generated = true
			function.analysis_root = true
		}
	}

	slice.sort_by(tests[:], test_entry_less)
	inject_assertion_locations(module, sources, symbols, testing_pkg)
	append_runner(module, sources, diagnostics, symbols, testing_pkg, error_type, tests[:])
	return true
}