package shader

import "core:mem"
import "core:strings"

// Public API for the Luma shader compiler

Target :: enum {
	GLSL_450,
	GLSL_450_OPENGL,
	HLSL_SM6,
	MSL_2_4,
	WGSL,
	SPIR_V,
	DXBC_SM5,
}

Opt_Level :: enum {
	None,
	Basic,
	Aggressive,
}

Compile_Options :: struct {
	target:       Target,
	opt_level:    Opt_Level,
	stage:        Shader_Stage,
	entry:        string,
	debug:        bool,
	validate:     bool,
	include_dirs: []string,
	file_reader:  File_Reader,

	// HLSL-specific (only used when target is HLSL_SM6)
	hlsl_cbuffer_push_constants: bool, // emit push constants as cbuffer (for D3D11/D3D12)
	hlsl_push_constant_slot:    int,   // register(b<N>) when hlsl_cbuffer_push_constants is true
	hlsl_omit_register_spaces:  bool,  // omit ", spaceN" from register() annotations
}

Compile_Result :: struct {
	success:     bool,
	output:      []u8,
	diagnostics: []Diagnostic,
	reflection:  Reflect_Info,
	allocator:   mem.Allocator,
}

// Primary entry point — source in, compiled output out.
compile :: proc(source: string, options: Compile_Options, file := "", allocator := context.allocator) -> Compile_Result {
	scratch: mem.Dynamic_Arena
	mem.dynamic_arena_init(&scratch, allocator, allocator, alignment = 64)
	defer mem.dynamic_arena_destroy(&scratch)

	scratch_allocator := mem.dynamic_arena_allocator(&scratch)
	old_allocator := context.allocator
	old_temp_allocator := context.temp_allocator
	context.allocator = scratch_allocator
	context.temp_allocator = scratch_allocator
	defer {
		context.allocator = old_allocator
		context.temp_allocator = old_temp_allocator
	}

	scratch_result := compile_unowned(source, options, file)

	context.allocator = old_allocator
	context.temp_allocator = old_temp_allocator
	return clone_compile_result(scratch_result, allocator)
}

destroy_compile_result :: proc(result: ^Compile_Result) {
	if result == nil do return
	if result.allocator.procedure == nil {
		result^ = {}
		return
	}

	if len(result.output) > 0 {
		delete(result.output, result.allocator)
	}

	for &d in result.diagnostics {
		if len(d.message) > 0 {
			delete(d.message, result.allocator)
		}
		if len(d.span.file) > 0 {
			delete(d.span.file, result.allocator)
		}
	}
	if len(result.diagnostics) > 0 {
		delete(result.diagnostics, result.allocator)
	}

	destroy_reflect_info(&result.reflection, result.allocator)

	result^ = {}
}

@(private)
compile_unowned :: proc(source: string, options: Compile_Options, file := "") -> Compile_Result {
	all_diags := make([dynamic]Diagnostic)

	// Preprocess
	pp := preprocessor_init(options.include_dirs, options.file_reader)
	processed_source, pp_diags := preprocess(&pp, source, file)
	for d in pp_diags do append(&all_diags, d)

	if has_errors(all_diags) {
		return Compile_Result{
			success     = false,
			diagnostics = all_diags[:],
		}
	}

	// Lex
	lex := lexer_init(processed_source, file)
	tokens, lex_diags := tokenize(&lex)
	for d in lex_diags do append(&all_diags, d)

	if has_errors(all_diags) {
		return Compile_Result{
			success     = false,
			diagnostics = all_diags[:],
		}
	}

	// Parse
	parser := parser_init(tokens, file)
	mod, parse_diags := parse_module(&parser)
	for d in parse_diags do append(&all_diags, d)

	if has_errors(all_diags) {
		return Compile_Result{
			success     = false,
			diagnostics = all_diags[:],
		}
	}

	// Desugar inline entry points
	desugar_module(mod)

	// Semantic analysis
	sema := sema_init()
	sema_diags := check_module(&sema, mod)
	for d in sema_diags do append(&all_diags, d)

	if has_errors(all_diags) {
		return Compile_Result{
			success     = false,
			diagnostics = all_diags[:],
		}
	}

	// Lower to IR
	ir_module, ir_diags := ir_build_module(mod, &sema)
	for d in ir_diags do append(&all_diags, d)

	if has_errors(all_diags) {
		return Compile_Result{
			success     = false,
			diagnostics = all_diags[:],
		}
	}

	// Optimize
	optimize(&ir_module, options.opt_level)

	// Filter entry points if --entry specified
	if options.entry != "" {
		filtered := make([dynamic]IR_Function)
		for fn in ir_module.functions {
			if !fn.is_entry || fn.name == options.entry {
				append(&filtered, fn)
			}
		}
		ir_module.functions = filtered
	}

	reflection := ir_reflect(&ir_module)

	// Emit
	output_bytes: []u8
	emit_diags: []Diagnostic

	switch options.target {
	case .GLSL_450:
		str: string
		str, emit_diags = emit_glsl(&ir_module)
		output_bytes = transmute([]u8)str
	case .GLSL_450_OPENGL:
		str: string
		str, emit_diags = emit_glsl_opengl(&ir_module)
		output_bytes = transmute([]u8)str
	case .WGSL:
		str: string
		str, emit_diags = emit_wgsl(&ir_module)
		output_bytes = transmute([]u8)str
	case .HLSL_SM6:
		str: string
		opts := options
		str, emit_diags = emit_hlsl(&ir_module, &opts)
		output_bytes = transmute([]u8)str
	case .MSL_2_4:
		str: string
		str, emit_diags = emit_msl(&ir_module)
		output_bytes = transmute([]u8)str
	case .SPIR_V:
		output_bytes, emit_diags = emit_spirv(&ir_module, options.debug, file)
	case .DXBC_SM5:
		output_bytes, emit_diags = emit_dxbc(&ir_module)
	}

	for d in emit_diags do append(&all_diags, d)

	return Compile_Result{
		success     = !has_errors(all_diags),
		output      = output_bytes,
		diagnostics = all_diags[:],
		reflection  = reflection,
	}
}

@(private)
clone_compile_result :: proc(result: Compile_Result, allocator: mem.Allocator) -> Compile_Result {
	return {
		success     = result.success,
		output      = clone_bytes(result.output, allocator),
		diagnostics = clone_diagnostics(result.diagnostics, allocator),
		reflection  = clone_reflect_info(result.reflection, allocator),
		allocator   = allocator,
	}
}

@(private)
clone_bytes :: proc(bytes: []u8, allocator: mem.Allocator) -> []u8 {
	if len(bytes) == 0 do return nil
	out := make([]u8, len(bytes), allocator)
	copy(out, bytes)
	return out
}

@(private)
clone_string_owned :: proc(value: string, allocator: mem.Allocator) -> string {
	if len(value) == 0 do return ""
	return strings.clone(value, allocator)
}

@(private)
clone_source_span :: proc(span: Source_Span, allocator: mem.Allocator) -> Source_Span {
	return {
		file       = clone_string_owned(span.file, allocator),
		line_start = span.line_start,
		col_start  = span.col_start,
		line_end   = span.line_end,
		col_end    = span.col_end,
	}
}

@(private)
clone_diagnostics :: proc(diagnostics: []Diagnostic, allocator: mem.Allocator) -> []Diagnostic {
	if len(diagnostics) == 0 do return nil
	out := make([]Diagnostic, len(diagnostics), allocator)
	for d, i in diagnostics {
		out[i] = {
			level   = d.level,
			message = clone_string_owned(d.message, allocator),
			span    = clone_source_span(d.span, allocator),
		}
	}
	return out
}

@(private)
clone_reflect_io :: proc(values: [dynamic]Reflect_IO, allocator: mem.Allocator) -> [dynamic]Reflect_IO {
	out := make([dynamic]Reflect_IO, 0, len(values), allocator)
	for value in values {
		append(&out, Reflect_IO{
			name     = clone_string_owned(value.name, allocator),
			type_str = clone_string_owned(value.type_str, allocator),
			location = value.location,
			builtin  = clone_string_owned(value.builtin, allocator),
		})
	}
	return out
}

@(private)
clone_reflect_struct_fields :: proc(values: [dynamic]Reflect_Struct_Field, allocator: mem.Allocator) -> [dynamic]Reflect_Struct_Field {
	out := make([dynamic]Reflect_Struct_Field, 0, len(values), allocator)
	for value in values {
		append(&out, Reflect_Struct_Field{
			name     = clone_string_owned(value.name, allocator),
			type_str = clone_string_owned(value.type_str, allocator),
			offset   = value.offset,
			size     = value.size,
		})
	}
	return out
}

@(private)
clone_reflect_info :: proc(info: Reflect_Info, allocator: mem.Allocator) -> Reflect_Info {
	out := Reflect_Info{
		entry_points   = make([dynamic]Reflect_Entry, 0, len(info.entry_points), allocator),
		structs        = make([dynamic]Reflect_Struct, 0, len(info.structs), allocator),
		bindings       = make([dynamic]Reflect_Binding, 0, len(info.bindings), allocator),
		spec_constants = make([dynamic]Reflect_Spec_Constant, 0, len(info.spec_constants), allocator),
	}

	for entry in info.entry_points {
		append(&out.entry_points, Reflect_Entry{
			name      = clone_string_owned(entry.name, allocator),
			stage     = entry.stage,
			inputs    = clone_reflect_io(entry.inputs, allocator),
			outputs   = clone_reflect_io(entry.outputs, allocator),
			workgroup = entry.workgroup,
		})
	}

	for st in info.structs {
		append(&out.structs, Reflect_Struct{
			name   = clone_string_owned(st.name, allocator),
			fields = clone_reflect_struct_fields(st.fields, allocator),
			size   = st.size,
		})
	}

	for binding in info.bindings {
		append(&out.bindings, Reflect_Binding{
			name        = clone_string_owned(binding.name, allocator),
			group       = binding.group,
			binding     = binding.binding,
			kind        = binding.kind,
			struct_name = clone_string_owned(binding.struct_name, allocator),
			size        = binding.size,
			combined_id = binding.combined_id,
		})
	}

	for spec in info.spec_constants {
		append(&out.spec_constants, Reflect_Spec_Constant{
			name          = clone_string_owned(spec.name, allocator),
			spec_id       = spec.spec_id,
			type_str      = clone_string_owned(spec.type_str, allocator),
			default_value = clone_string_owned(spec.default_value, allocator),
		})
	}

	return out
}

@(private)
destroy_reflect_info :: proc(info: ^Reflect_Info, allocator: mem.Allocator) {
	for &entry in info.entry_points {
		if len(entry.name) > 0 do delete(entry.name, allocator)
		for &input in entry.inputs {
			if len(input.name) > 0 do delete(input.name, allocator)
			if len(input.type_str) > 0 do delete(input.type_str, allocator)
			if len(input.builtin) > 0 do delete(input.builtin, allocator)
		}
		if len(entry.inputs) > 0 do delete(entry.inputs)
		for &output in entry.outputs {
			if len(output.name) > 0 do delete(output.name, allocator)
			if len(output.type_str) > 0 do delete(output.type_str, allocator)
			if len(output.builtin) > 0 do delete(output.builtin, allocator)
		}
		if len(entry.outputs) > 0 do delete(entry.outputs)
	}
	if len(info.entry_points) > 0 do delete(info.entry_points)

	for &st in info.structs {
		if len(st.name) > 0 do delete(st.name, allocator)
		for &field in st.fields {
			if len(field.name) > 0 do delete(field.name, allocator)
			if len(field.type_str) > 0 do delete(field.type_str, allocator)
		}
		if len(st.fields) > 0 do delete(st.fields)
	}
	if len(info.structs) > 0 do delete(info.structs)

	for &binding in info.bindings {
		if len(binding.name) > 0 do delete(binding.name, allocator)
		if len(binding.struct_name) > 0 do delete(binding.struct_name, allocator)
	}
	if len(info.bindings) > 0 do delete(info.bindings)

	for &spec in info.spec_constants {
		if len(spec.name) > 0 do delete(spec.name, allocator)
		if len(spec.type_str) > 0 do delete(spec.type_str, allocator)
		if len(spec.default_value) > 0 do delete(spec.default_value, allocator)
	}
	if len(info.spec_constants) > 0 do delete(info.spec_constants)

	info^ = {}
}

// Pipeline stage APIs for advanced usage

compile_parse :: proc(source: string, file := "") -> (^Ast_Module, []Diagnostic) {
	lex := lexer_init(source, file)
	tokens, lex_diags := tokenize(&lex)

	all_diags := make([dynamic]Diagnostic)
	for d in lex_diags do append(&all_diags, d)

	if has_errors(all_diags) {
		return nil, all_diags[:]
	}

	parser := parser_init(tokens, file)
	return parse_module(&parser)
}

// Lower source to IR (preprocess → lex → parse → sema → IR)
compile_lower :: proc(source: string, file := "", opt: Opt_Level = .None) -> (IR_Module, []Diagnostic) {
	all_diags := make([dynamic]Diagnostic)

	lex := lexer_init(source, file)
	tokens, lex_diags := tokenize(&lex)
	for d in lex_diags do append(&all_diags, d)
	if has_errors(all_diags) do return {}, all_diags[:]

	parser := parser_init(tokens, file)
	mod, parse_diags := parse_module(&parser)
	for d in parse_diags do append(&all_diags, d)
	if has_errors(all_diags) do return {}, all_diags[:]

	desugar_module(mod)

	sema := sema_init()
	sema_diags := check_module(&sema, mod)
	for d in sema_diags do append(&all_diags, d)
	if has_errors(all_diags) do return {}, all_diags[:]

	ir_module, ir_diags := ir_build_module(mod, &sema)
	for d in ir_diags do append(&all_diags, d)

	if !has_errors(all_diags) {
		optimize(&ir_module, opt)
	}

	return ir_module, all_diags[:]
}

// Extract reflection info from source
compile_reflect :: proc(source: string, file := "") -> (Reflect_Info, []Diagnostic) {
	ir_module, diags := compile_lower(source, file)
	if has_errors_slice(diags) do return {}, diags
	return ir_reflect(&ir_module), diags
}

// Extract SPIR-V-specific reflection info from source
compile_reflect_spirv :: proc(source: string, file := "") -> (Reflect_SPIRV_Info, []Diagnostic) {
	ir_module, diags := compile_lower(source, file)
	if has_errors_slice(diags) do return {}, diags
	return ir_reflect_spirv(&ir_module), diags
}

// Parse and desugar source, returning the AST for dump-ast
compile_dump_ast :: proc(source: string, file := "") -> (^Ast_Module, []Diagnostic) {
	all_diags := make([dynamic]Diagnostic)

	lex := lexer_init(source, file)
	tokens, lex_diags := tokenize(&lex)
	for d in lex_diags do append(&all_diags, d)
	if has_errors(all_diags) do return nil, all_diags[:]

	parser := parser_init(tokens, file)
	mod, parse_diags := parse_module(&parser)
	for d in parse_diags do append(&all_diags, d)
	if has_errors(all_diags) do return nil, all_diags[:]

	desugar_module(mod)
	return mod, all_diags[:]
}

compile_check :: proc(source: string, file := "") -> (^Ast_Module, ^Sema, []Diagnostic) {
	all_diags := make([dynamic]Diagnostic)

	lex := lexer_init(source, file)
	tokens, lex_diags := tokenize(&lex)
	for d in lex_diags do append(&all_diags, d)
	if has_errors(all_diags) do return nil, nil, all_diags[:]

	parser := parser_init(tokens, file)
	mod, parse_diags := parse_module(&parser)
	for d in parse_diags do append(&all_diags, d)
	if has_errors(all_diags) do return nil, nil, all_diags[:]

	desugar_module(mod)

	sema := new(Sema)
	sema^ = sema_init()
	sema_diags := check_module(sema, mod)
	for d in sema_diags do append(&all_diags, d)

	return mod, sema, all_diags[:]
}