package compiler

import "./ast"
import "./backend"
import "./cimport"
import "./checker"
import "./llvm"
import "./linker"
import "./loader"
import "./lower"
import "./opt"
import "./source"
import "./symbol"
import "./target"
import "core:fmt"
import vmem "core:mem/virtual"
import "core:os"
import "core:os/os2"
import "core:strings"

// write_escaped_c_string writes value into builder with `\` and `"` escaped so
// it is safe to embed inside a C string literal (e.g. an `#include "..."`).
write_escaped_c_string :: proc(builder: ^strings.Builder, value: string) {
	for b in transmute([]byte)value {
		if b == '\\' || b == '"' {
			strings.write_byte(builder, '\\')
		}
		strings.write_byte(builder, b)
	}
}

compile_package :: proc(
	input_path, output_path: string,
	link_arguments: []linker.Argument = nil,
	selected := target.DEFAULT,
	c_options := cimport.Options{},
	project_root := "",
	mode := ast.Compile_Mode.Executable,
) -> int {
	sources := source.init_store()
	defer source.destroy_store(&sources)
	diagnostics := source.init_store_diagnostics(&sources)
	defer source.destroy_diagnostics(&diagnostics)
	symbols := symbol.init_table()
	defer symbol.destroy_table(&symbols)

	lexer_arena: vmem.Arena
	if err := vmem.arena_init_growing(&lexer_arena); err != nil {
		fmt.eprintln("failed to initialize lexer arena:", err)
		return 2
	}
	defer vmem.arena_destroy(&lexer_arena)
	parser_arena: vmem.Arena
	if err := vmem.arena_init_growing(&parser_arena); err != nil {
		fmt.eprintln("failed to initialize parser arena:", err)
		return 2
	}
	defer vmem.arena_destroy(&parser_arena)
	checker_arena: vmem.Arena
	if err := vmem.arena_init_growing(&checker_arena); err != nil {
		fmt.eprintln("failed to initialize checker arena:", err)
		return 2
	}
	defer vmem.arena_destroy(&checker_arena)
	lower_arena: vmem.Arena
	if err := vmem.arena_init_growing(&lower_arena); err != nil {
		fmt.eprintln("failed to initialize lowering arena:", err)
		return 2
	}
	defer vmem.arena_destroy(&lower_arena)

	ast_module, loaded := loader.load(
		input_path,
		&sources,
		&diagnostics,
		&symbols,
		vmem.arena_allocator(&lexer_arena),
		vmem.arena_allocator(&parser_arena),
		c_options,
		selected,
		project_root if len(project_root) > 0 else input_path,
		mode,
	)
	if !loaded {
		source.print_all(&diagnostics)
		fmt.eprintln("failed to load root package directory:", input_path)
		return 2
	}
	effective_root := project_root if len(project_root) > 0 else input_path
	if !prepare_tests(&ast_module, &sources, &diagnostics, &symbols, mode, effective_root) {
		source.print_all(&diagnostics)
		return 1
	}

	// Generated C trampolines (for `static inline` imports) must be compiled and
	// linked with the program. Write them out and add the source as a link input
	// before the loader's arena (which owns the trampoline strings) is freed.
	trampoline_path := ""
	effective_link_arguments := link_arguments
	owns_arguments := false
	// Register the path defer first so it runs last (LIFO): the augmented
	// argument slice, whose Input value aliases trampoline_path, is freed before
	// the path string it points at.
	defer if len(trampoline_path) > 0 {
		_ = os.remove(trampoline_path)
		delete(trampoline_path)
	}
	defer if owns_arguments {
		delete(effective_link_arguments)
	}
	if len(ast_module.c_trampolines) > 0 {
		builder := strings.builder_make()
		defer strings.builder_destroy(&builder)
		seen_headers: map[string]bool
		defer delete(seen_headers)
		for trampoline in ast_module.c_trampolines {
			if seen_headers[trampoline.header] {
				continue
			}
			seen_headers[trampoline.header] = true
			strings.write_string(&builder, "#include \"")
			write_escaped_c_string(&builder, trampoline.header)
			strings.write_string(&builder, "\"\n")
		}
		for trampoline in ast_module.c_trampolines {
			strings.write_string(&builder, trampoline.source)
		}
		// Owned (stable allocator); freed in the path defer above after backend
		// compilation, which is the only consumer, has run.
		trampoline_path = fmt.aprintf("%s.brolang-trampolines-%d.c", output_path, os2.get_pid())
		if err := os.write_entire_file_or_err(trampoline_path, transmute([]byte)strings.to_string(builder)); err != nil {
			fmt.eprintln("failed to write C trampoline source:", err)
			return 2
		}
		augmented := make([]linker.Argument, len(link_arguments)+1)
		copy(augmented, link_arguments)
		augmented[len(link_arguments)] = linker.Argument{kind=.Input, value=trampoline_path}
		effective_link_arguments = augmented
		owns_arguments = true
	}

	vmem.arena_free_all(&lexer_arena)
	hir_module := checker.check(&ast_module, &diagnostics, &symbols, selected, vmem.arena_allocator(&checker_arena))
	vmem.arena_free_all(&parser_arena)
	ir_module := lower.lower(&hir_module, vmem.arena_allocator(&lower_arena))
	vmem.arena_free_all(&checker_arena)
	opt.run(&ir_module)

	llvm_text := llvm.emit(&ir_module, &diagnostics, &symbols)
	defer delete(llvm_text)
	vmem.arena_free_all(&lower_arena)
	llvm_path := fmt.tprintf("%s.brolang-%d.ll", output_path, os2.get_pid())
	defer _ = os.remove(llvm_path)
	if err := os.write_entire_file_or_err(llvm_path, transmute([]byte)llvm_text); err != nil {
		fmt.eprintln("failed to write temporary LLVM IR:", err)
		return 2
	}

	source.print_all(&diagnostics)
	if !backend.compile(llvm_path, output_path, effective_link_arguments, selected, c_options) {
		return 2
	}
	if len(diagnostics.items) > 0 {
		return 1
	}
	return 0
}