package main

import "core:fmt"
import "core:mem"
import "core:os"
import "core:strings"
import shader "../../shader"

Scratch_Scope :: struct {
	arena:              mem.Dynamic_Arena,
	old_allocator:      mem.Allocator,
	old_temp_allocator: mem.Allocator,
}

scratch_begin :: proc(scope: ^Scratch_Scope) {
	scope.old_allocator = context.allocator
	scope.old_temp_allocator = context.temp_allocator
	mem.dynamic_arena_init(&scope.arena, context.allocator, context.allocator, alignment = 64)
	allocator := mem.dynamic_arena_allocator(&scope.arena)
	context.allocator = allocator
	context.temp_allocator = allocator
}

scratch_end :: proc(scope: ^Scratch_Scope) {
	context.allocator = scope.old_allocator
	context.temp_allocator = scope.old_temp_allocator
	mem.dynamic_arena_destroy(&scope.arena)
	scope^ = {}
}

main :: proc() {
	args := os.args[1:]

	if len(args) == 0 {
		print_usage()
		os.exit(1)
	}

	command := args[0]
	rest := args[1:]

	switch command {
	case "compile":
		cmd_compile(rest)
	case "check":
		cmd_check(rest)
	case "dump-ast":
		cmd_dump_ast(rest)
	case "dump-ir":
		cmd_dump_ir(rest)
	case "reflect":
		cmd_reflect(rest)
	case "help", "--help", "-h":
		print_usage()
	case:
		fmt.eprintfln("unknown command: %s", command)
		print_usage()
		os.exit(1)
	}
}

cmd_compile :: proc(args: []string) {
	input_file := ""
	output_file := ""
	target := shader.Target.GLSL_450
	include_dirs := make([dynamic]string)
	debug := false
	entry := ""

	i := 0
	for i < len(args) {
		arg := args[i]
		if arg == "-o" && i + 1 < len(args) {
			i += 1
			output_file = args[i]
		} else if arg == "--debug" {
			debug = true
		} else if strings.has_prefix(arg, "--entry=") {
			entry = arg[len("--entry="):]
		} else if strings.has_prefix(arg, "--target=") {
			target_str := arg[len("--target="):]
			switch target_str {
			case "glsl":        target = .GLSL_450
			case "glsl-opengl": target = .GLSL_450_OPENGL
			case "hlsl":        target = .HLSL_SM6
			case "msl":         target = .MSL_2_4
			case "wgsl":        target = .WGSL
			case "spirv":       target = .SPIR_V
			case "dxbc":        target = .DXBC_SM5
			case:
				fmt.eprintfln("unknown target: %s", target_str)
				os.exit(1)
			}
		} else if strings.has_prefix(arg, "--include-dir=") {
			append(&include_dirs, arg[len("--include-dir="):])
		} else if len(arg) > 0 && arg[0] != '-' {
			input_file = arg
		} else {
			fmt.eprintfln("unknown option: %s", arg)
			os.exit(1)
		}
		i += 1
	}

	if input_file == "" {
		fmt.eprintln("error: no input file specified")
		os.exit(1)
	}

	source_data, read_err := os.read_entire_file(input_file, context.allocator)
	if read_err != nil {
		fmt.eprintfln("error: could not read file: %s", input_file)
		os.exit(1)
	}
	source := string(source_data)
	defer delete(source_data)

	options := shader.Compile_Options{
		target       = target,
		include_dirs = include_dirs[:],
		debug        = debug,
		entry        = entry,
	}

	result := shader.compile(source, options, input_file)
	defer shader.destroy_compile_result(&result)

	// Print diagnostics
	for d in result.diagnostics {
		msg := shader.format_diagnostic_with_source(d, source)
		fmt.eprintln(msg)
		delete(msg)
	}

	if !result.success {
		os.exit(1)
	}

	if output_file != "" {
		ok := os.write_entire_file(output_file, result.output)
		if ok != nil {
			fmt.eprintfln("error: could not write file: %s", output_file)
			os.exit(1)
		}
		fmt.printfln("wrote %s", output_file)
	} else {
		fmt.print(string(result.output))
	}
}

cmd_check :: proc(args: []string) {
	input_file := ""
	include_dirs := make([dynamic]string)

	for arg in args {
		if strings.has_prefix(arg, "--include-dir=") {
			append(&include_dirs, arg[len("--include-dir="):])
		} else if len(arg) > 0 && arg[0] != '-' {
			input_file = arg
		}
	}

	if input_file == "" {
		fmt.eprintln("error: no input file specified")
		os.exit(1)
	}

	source_data, read_err := os.read_entire_file(input_file, context.allocator)
	if read_err != nil {
		fmt.eprintfln("error: could not read file: %s", input_file)
		os.exit(1)
	}
	source := string(source_data)
	defer delete(source_data)

	scratch: Scratch_Scope
	scratch_begin(&scratch)
	defer scratch_end(&scratch)

	_, _, diags := shader.compile_check(source, input_file)

	has_errors := false
	for d in diags {
		fmt.eprintln(shader.format_diagnostic_with_source(d, source))
		if d.level == .Error do has_errors = true
	}

	if has_errors {
		os.exit(1)
	} else {
		fmt.printfln("%s: ok", input_file)
	}
}

cmd_dump_ast :: proc(args: []string) {
	input_file := ""
	for arg in args {
		if len(arg) > 0 && arg[0] != '-' {
			input_file = arg
		}
	}

	if input_file == "" {
		fmt.eprintln("error: no input file specified")
		os.exit(1)
	}

	source_data, read_err := os.read_entire_file(input_file, context.allocator)
	if read_err != nil {
		fmt.eprintfln("error: could not read file: %s", input_file)
		os.exit(1)
	}
	source := string(source_data)
	defer delete(source_data)

	scratch: Scratch_Scope
	scratch_begin(&scratch)
	defer scratch_end(&scratch)

	mod, diags := shader.compile_dump_ast(source, input_file)

	has_errors := false
	for d in diags {
		fmt.eprintln(shader.format_diagnostic_with_source(d, source))
		if d.level == .Error do has_errors = true
	}
	if has_errors do os.exit(1)

	fmt.print(shader.ast_to_string(mod))
}

cmd_dump_ir :: proc(args: []string) {
	input_file := ""
	for arg in args {
		if len(arg) > 0 && arg[0] != '-' {
			input_file = arg
		}
	}

	if input_file == "" {
		fmt.eprintln("error: no input file specified")
		os.exit(1)
	}

	source_data, read_err := os.read_entire_file(input_file, context.allocator)
	if read_err != nil {
		fmt.eprintfln("error: could not read file: %s", input_file)
		os.exit(1)
	}
	source := string(source_data)
	defer delete(source_data)

	scratch: Scratch_Scope
	scratch_begin(&scratch)
	defer scratch_end(&scratch)

	ir_module, diags := shader.compile_lower(source, input_file)

	has_errors := false
	for d in diags {
		fmt.eprintln(shader.format_diagnostic_with_source(d, source))
		if d.level == .Error do has_errors = true
	}
	if has_errors do os.exit(1)

	fmt.print(shader.ir_to_string(&ir_module))
}

cmd_reflect :: proc(args: []string) {
	input_file := ""
	spirv_mode := false
	for arg in args {
		if arg == "--spirv" {
			spirv_mode = true
		} else if len(arg) > 0 && arg[0] != '-' {
			input_file = arg
		}
	}

	if input_file == "" {
		fmt.eprintln("error: no input file specified")
		os.exit(1)
	}

	source_data, read_err := os.read_entire_file(input_file, context.allocator)
	if read_err != nil {
		fmt.eprintfln("error: could not read file: %s", input_file)
		os.exit(1)
	}
	source := string(source_data)
	defer delete(source_data)

	if spirv_mode {
		scratch: Scratch_Scope
		scratch_begin(&scratch)
		defer scratch_end(&scratch)

		info, diags := shader.compile_reflect_spirv(source, input_file)
		has_errors := false
		for d in diags {
			fmt.eprintln(shader.format_diagnostic_with_source(d, source))
			if d.level == .Error do has_errors = true
		}
		if has_errors do os.exit(1)
		fmt.print(shader.reflect_spirv_to_json(info))
	} else {
		scratch: Scratch_Scope
		scratch_begin(&scratch)
		defer scratch_end(&scratch)

		info, diags := shader.compile_reflect(source, input_file)
		has_errors := false
		for d in diags {
			fmt.eprintln(shader.format_diagnostic_with_source(d, source))
			if d.level == .Error do has_errors = true
		}
		if has_errors do os.exit(1)
		fmt.print(shader.reflect_to_json(info))
	}
}

print_usage :: proc() {
	fmt.println("Usage: luma <command> [options] <file>")
	fmt.println("")
	fmt.println("Commands:")
	fmt.println("  compile   Compile a .luma shader file")
	fmt.println("  check     Parse and typecheck only")
	fmt.println("  dump-ast  Dump abstract syntax tree")
	fmt.println("  dump-ir   Dump intermediate representation")
	fmt.println("  reflect   Emit binding reflection JSON")
	fmt.println("  help      Show this help message")
	fmt.println("")
	fmt.println("Compile options:")
	fmt.println("  --target=<target>       Output target: glsl, hlsl, msl, wgsl, spirv, dxbc (default: glsl)")
	fmt.println("  --include-dir=<path>    Add include search path")
	fmt.println("  --entry=<name>          Emit only the named entry point (default: all)")
	fmt.println("  --debug                 Emit debug info (SPIR-V: OpLine/OpSource)")
	fmt.println("  -o <file>               Output file (default: stdout)")
	fmt.println("")
	fmt.println("Reflect options:")
	fmt.println("  --spirv                 Emit SPIR-V-specific reflection (descriptor sets, layout)")
}