package llvm

import "../ir"
import "../source"
import "../symbol"
import "../types"
import "core:fmt"
import "core:mem"
import "core:strings"

Trap_Message :: struct {
	text: string,
}

Emitter :: struct {
	module:      ^ir.Module,
	diagnostics: ^source.Diagnostics,
	symbols:     ^symbol.Table,
	builder:     strings.Builder,
	messages:    [dynamic]Trap_Message,
	allocator:   mem.Allocator,
}

llvm_type :: proc(value: types.Type) -> string {
	if value.kind == .Void {
		return "void"
	}
	switch value.bits {
	case 8:  return "i8"
	case 16: return "i16"
	case 32: return "i32"
	case:    return "i64"
	}
}

function_result_type :: proc(function: ir.Function) -> string {
	if function.is_main {
		return "i32"
	}
	return llvm_type(function.result)
}

write_operand :: proc(builder: ^strings.Builder, instructions: []ir.Instruction, value_id: int) {
	if value_id < 0 || value_id >= len(instructions) {
		fmt.sbprintf(builder, "-6148914691236517206")
		return
	}
	value := instructions[value_id]
	if value.op == .Const {
		fmt.sbprintf(builder, "%d", value.integer)
	} else {
		fmt.sbprintf(builder, "%%v%d", value_id)
	}
}

register_message :: proc(emitter: ^Emitter, text: string) -> int {
	id := len(emitter.messages)
	cloned := fmt.aprintf("%s\n", text, allocator=emitter.allocator)
	append(&emitter.messages, Trap_Message{text=cloned})
	return id
}

diagnostic_message :: proc(emitter: ^Emitter, diagnostic: int, span: source.Span, fallback: string) -> int {
	if diagnostic >= 0 && diagnostic < len(emitter.diagnostics.items) {
		message := source.format(emitter.diagnostics, diagnostic, emitter.allocator)
		id := register_message(emitter, message)
		delete(message, emitter.allocator)
		return id
	}
	source_file := source.source_for_span(emitter.diagnostics, span)
	if source_file == nil {
		return register_message(emitter, fallback)
	}
	line, column := source.line_and_column(source_file, span.start)
	message := fmt.aprintf(
		"%s:%d:%d: runtime trap: %s",
		source_file.path,
		line,
		column,
		fallback,
		allocator=emitter.allocator,
	)
	id := register_message(emitter, message)
	delete(message, emitter.allocator)
	return id
}

emit_trap_call :: proc(emitter: ^Emitter, message_id: int) {
	message := emitter.messages[message_id]
	fmt.sbprintf(
		&emitter.builder,
		"  call void @bro.trap(ptr @bro.msg.%d, i64 %d)\n",
		message_id,
		len(message.text),
	)
}

emit_call_args :: proc(builder: ^strings.Builder, instructions: []ir.Instruction, args: []int) {
	for arg, index in args {
		if index > 0 {
			strings.write_string(builder, ", ")
		}
		fmt.sbprintf(builder, "%s ", llvm_type(instructions[arg].type))
		write_operand(builder, instructions, arg)
	}
}

emit_instruction_stream :: proc(
	emitter: ^Emitter,
	instructions: []ir.Instruction,
	function: ir.Function,
	global_initializer := false,
) -> int {
	return_value := -1
	after_return := false
	for instruction, instruction_id in instructions {
		if after_return {
			fmt.sbprintf(&emitter.builder, "recover_after_return_%d:\n", instruction_id)
			after_return = false
		}
		switch instruction.op {
		case .Param, .Const:
		case .Load_Global:
			if instruction.target < 0 || instruction.target >= len(emitter.module.globals) {
				message := diagnostic_message(emitter, -1, instruction.span, "invalid global reference")
				emit_trap_call(emitter, message)
				continue
			}
			global := emitter.module.globals[instruction.target]
			if global.is_static {
				fmt.sbprintf(
					&emitter.builder,
					"  %%v%d = load %s, ptr @bro.g.%d\n",
					instruction_id,
					llvm_type(global.type),
					instruction.target,
				)
			} else {
				fmt.sbprintf(
					&emitter.builder,
					"  %%v%d = call %s @bro.get.%d()\n",
					instruction_id,
					llvm_type(global.type),
					instruction.target,
				)
			}
		case .Alloca:
			fmt.sbprintf(&emitter.builder, "  %%v%d = alloca %s\n", instruction_id, llvm_type(instruction.type))
		case .Load:
			fmt.sbprintf(
				&emitter.builder,
				"  %%v%d = load %s, ptr %%v%d\n",
				instruction_id,
				llvm_type(instruction.type),
				instruction.a,
			)
		case .Store:
			fmt.sbprintf(&emitter.builder, "  store %s ", llvm_type(instruction.type))
			write_operand(&emitter.builder, instructions, instruction.b)
			fmt.sbprintf(&emitter.builder, ", ptr %%v%d\n", instruction.a)
		case .Widen:
			from_type := instructions[instruction.a].type
			fmt.sbprintf(&emitter.builder, "  %%v%d = sext %s ", instruction_id, llvm_type(from_type))
			write_operand(&emitter.builder, instructions, instruction.a)
			fmt.sbprintf(&emitter.builder, " to %s\n", llvm_type(instruction.type))
		case .Add_Checked:
			type_name := llvm_type(instruction.type)
			fmt.sbprintf(&emitter.builder, "  %%pair%d = call ", instruction_id)
			strings.write_string(&emitter.builder, "{ ")
			fmt.sbprintf(&emitter.builder, "%s, i1 } @llvm.sadd.with.overflow.%s(%s ", type_name, type_name, type_name)
			write_operand(&emitter.builder, instructions, instruction.a)
			fmt.sbprintf(&emitter.builder, ", %s ", type_name)
			write_operand(&emitter.builder, instructions, instruction.b)
			fmt.sbprintf(&emitter.builder, ")\n")
			fmt.sbprintf(&emitter.builder, "  %%v%d = extractvalue ", instruction_id)
			strings.write_string(&emitter.builder, "{ ")
			fmt.sbprintf(&emitter.builder, "%s, i1 } %%pair%d, 0\n", type_name, instruction_id)
			fmt.sbprintf(&emitter.builder, "  %%overflow%d = extractvalue ", instruction_id)
			strings.write_string(&emitter.builder, "{ ")
			fmt.sbprintf(&emitter.builder, "%s, i1 } %%pair%d, 1\n", type_name, instruction_id)
			fmt.sbprintf(
				&emitter.builder,
				"  br i1 %%overflow%d, label %%overflow_trap%d, label %%overflow_continue%d\n",
				instruction_id,
				instruction_id,
				instruction_id,
			)
			fmt.sbprintf(&emitter.builder, "overflow_trap%d:\n", instruction_id)
			message := diagnostic_message(emitter, -1, instruction.span, "signed integer addition overflow")
			emit_trap_call(emitter, message)
			fmt.sbprintf(&emitter.builder, "  unreachable\noverflow_continue%d:\n", instruction_id)
		case .Call:
			if instruction.target < 0 || instruction.target >= len(emitter.module.functions) {
				message := diagnostic_message(emitter, -1, instruction.span, "invalid function specialization")
				emit_trap_call(emitter, message)
				continue
			}
			target := emitter.module.functions[instruction.target]
			if instruction.type.kind != .Void {
				fmt.sbprintf(&emitter.builder, "  %%v%d = ", instruction_id)
			} else {
				strings.write_string(&emitter.builder, "  ")
			}
			strings.write_string(&emitter.builder, "call ")
			if !target.c_abi {
				strings.write_string(&emitter.builder, "fastcc ")
			}
			fmt.sbprintf(&emitter.builder, "%s @%s(", function_result_type(target), target.link_name)
			emit_call_args(&emitter.builder, instructions, instruction.args)
			strings.write_string(&emitter.builder, ")\n")
		case .Trap:
			message := diagnostic_message(emitter, instruction.diagnostic, instruction.span, "invalid recovered source")
			emit_trap_call(emitter, message)
		case .Return:
			if global_initializer {
				return_value = instruction.a
				continue
			}
			fmt.sbprintf(&emitter.builder, "  ret %s ", function_result_type(function))
			write_operand(&emitter.builder, instructions, instruction.a)
			strings.write_string(&emitter.builder, "\n")
			after_return = true
		case .Return_Void:
			if global_initializer {
				continue
			}
			if function.is_main {
				strings.write_string(&emitter.builder, "  ret i32 0\n")
			} else {
				strings.write_string(&emitter.builder, "  ret void\n")
			}
			after_return = true
		}
	}
	return return_value
}

emit_globals :: proc(emitter: ^Emitter) {
	for global, global_id in emitter.module.globals {
		if global.is_static {
			fmt.sbprintf(
				&emitter.builder,
				"@bro.g.%d = internal constant %s %d\n",
				global_id,
				llvm_type(global.type),
				global.static_value,
			)
		} else {
			fmt.sbprintf(
				&emitter.builder,
				"@bro.g.%d = internal global %s 0\n@bro.gstate.%d = internal global i8 0\n",
				global_id,
				llvm_type(global.type),
				global_id,
			)
		}
	}
	strings.write_string(&emitter.builder, "\n")
}

emit_global_accessors :: proc(emitter: ^Emitter) {
	placeholder_function := ir.Function{result=types.I64}
	for global, global_id in emitter.module.globals {
		if global.is_static {
			continue
		}
		type_name := llvm_type(global.type)
		fmt.sbprintf(&emitter.builder, "define internal %s @bro.get.%d() ", type_name, global_id)
		strings.write_string(&emitter.builder, "{\nentry:\n")
		fmt.sbprintf(
			&emitter.builder,
			"  %%state = load i8, ptr @bro.gstate.%d\n  %%done = icmp eq i8 %%state, 2\n  br i1 %%done, label %%ready, label %%check\n",
			global_id,
		)
		strings.write_string(&emitter.builder, "check:\n  %visiting = icmp eq i8 %state, 1\n  br i1 %visiting, label %cycle, label %initialize\ncycle:\n")
		message_text := fmt.aprintf(
			"runtime trap: global initialization cycle involving '%s'",
			symbol.resolve(emitter.symbols, global.name),
			allocator=emitter.allocator,
		)
		message := register_message(emitter, message_text)
		delete(message_text, emitter.allocator)
		emit_trap_call(emitter, message)
		strings.write_string(&emitter.builder, "  unreachable\ninitialize:\n")
		fmt.sbprintf(&emitter.builder, "  store i8 1, ptr @bro.gstate.%d\n", global_id)
		placeholder_function.result = global.type
		value := emit_instruction_stream(emitter, global.initializer, placeholder_function, true)
		fmt.sbprintf(&emitter.builder, "  store %s ", type_name)
		write_operand(&emitter.builder, global.initializer, value)
		fmt.sbprintf(&emitter.builder, ", ptr @bro.g.%d\n", global_id)
		fmt.sbprintf(&emitter.builder, "  store i8 2, ptr @bro.gstate.%d\n", global_id)
		fmt.sbprintf(&emitter.builder, "  ret %s ", type_name)
		write_operand(&emitter.builder, global.initializer, value)
		strings.write_string(&emitter.builder, "\nready:\n")
		fmt.sbprintf(&emitter.builder, "  %%value = load %s, ptr @bro.g.%d\n  ret %s %%value\n}\n\n", type_name, global_id, type_name)
	}
}

emit_constructor :: proc(emitter: ^Emitter) {
	count := 0
	for global in emitter.module.globals {
		if !global.is_static && !global.problematic {
			count += 1
		}
	}
	if count == 0 {
		return
	}
	strings.write_string(
		&emitter.builder,
		"@llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr @bro.init, ptr null }]\n\n",
	)
	strings.write_string(&emitter.builder, "define internal void @bro.init() {\nentry:\n")
	for global, global_id in emitter.module.globals {
		if !global.is_static && !global.problematic {
			fmt.sbprintf(&emitter.builder, "  %%g%d = call %s @bro.get.%d()\n", global_id, llvm_type(global.type), global_id)
		}
	}
	strings.write_string(&emitter.builder, "  ret void\n}\n\n")
}

emit_functions :: proc(emitter: ^Emitter) {
	for function in emitter.module.functions {
		strings.write_string(&emitter.builder, "define ")
		if !function.c_abi {
			strings.write_string(&emitter.builder, "internal fastcc ")
		}
		fmt.sbprintf(&emitter.builder, "%s @%s(", function_result_type(function), function.link_name)
		for param_type, index in function.param_types {
			if index > 0 {
				strings.write_string(&emitter.builder, ", ")
			}
			fmt.sbprintf(&emitter.builder, "%s %%v%d", llvm_type(param_type), index)
		}
		strings.write_string(&emitter.builder, ") {\nentry:\n")
		_ = emit_instruction_stream(emitter, function.instructions, function)
		strings.write_string(&emitter.builder, "}\n\n")
	}
}

emit_escaped_bytes :: proc(builder: ^strings.Builder, text: string) {
	for value in transmute([]byte)text {
		if value >= 32 && value <= 126 && value != '\\' && value != '"' {
			strings.write_byte(builder, value)
		} else {
			fmt.sbprintf(builder, "\\%02X", value)
		}
	}
}

emit_messages :: proc(emitter: ^Emitter) {
	for message, message_id in emitter.messages {
		fmt.sbprintf(&emitter.builder, "@bro.msg.%d = private unnamed_addr constant [%d x i8] c\"", message_id, len(message.text))
		emit_escaped_bytes(&emitter.builder, message.text)
		strings.write_string(&emitter.builder, "\"\n")
	}
	strings.write_string(&emitter.builder, "\n")
}

emit_declarations :: proc(emitter: ^Emitter) {
	strings.write_string(&emitter.builder, "declare i64 @write(i32, ptr, i64)\ndeclare void @llvm.trap()\n")
	widths := [?]int{8, 16, 32, 64}
	for bits in widths {
		strings.write_string(&emitter.builder, "declare { i")
		fmt.sbprintf(&emitter.builder, "%d", bits)
		strings.write_string(&emitter.builder, ", i1 } @llvm.sadd.with.overflow.i")
		fmt.sbprintf(&emitter.builder, "%d(i%d, i%d)\n", bits, bits, bits)
	}
	strings.write_string(
		&emitter.builder,
		"\ndefine internal void @bro.trap(ptr %message, i64 %length) {\nentry:\n  %written = call i64 @write(i32 2, ptr %message, i64 %length)\n  call void @llvm.trap()\n  unreachable\n}\n\n",
	)
}

emit :: proc(
	module: ^ir.Module,
	diagnostics: ^source.Diagnostics,
	symbols: ^symbol.Table,
	allocator := context.allocator,
) -> string {
	emitter := Emitter{
		module=module,
		diagnostics=diagnostics,
		symbols=symbols,
		builder=strings.builder_make(allocator),
		allocator=allocator,
	}
	emitter.messages.allocator = allocator
	defer {
		for message in emitter.messages {
			delete(message.text, allocator)
		}
		delete(emitter.messages)
		strings.builder_destroy(&emitter.builder)
	}

	strings.write_string(&emitter.builder, "; generated by brolang\n\n")
	emit_globals(&emitter)
	emit_constructor(&emitter)
	emit_global_accessors(&emitter)
	emit_functions(&emitter)
	emit_messages(&emitter)
	emit_declarations(&emitter)
	return fmt.aprintf("%s", strings.to_string(emitter.builder), allocator=allocator)
}