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)
}

sentinel :: proc(value_type: types.Type) -> i64 {
	switch value_type.bits {
	case 8:  return -86
	case 16: return -21846
	case 32: return -1431655766
	case:    return -6148914691236517206
	}
}

valid_instruction :: proc(instructions: []ir.Instruction, instruction_id: ir.Instruction_Id) -> bool {
	return instruction_id != ir.INVALID_INSTRUCTION && int(instruction_id) < len(instructions)
}

valid_value :: proc(instructions: []ir.Instruction, value_id: ir.Instruction_Id, expected: types.Type) -> bool {
	if !valid_instruction(instructions, value_id) ||
	   !types.is_concrete_integer(expected) ||
	   !types.equal(instructions[value_id].type, expected) {
		return false
	}
	switch instructions[value_id].op {
	case .Param, .Const, .Load_Global, .Load, .Widen, .Neg_Checked, .Add_Checked, .Call:
		return true
	case .Alloca, .Store, .Trap, .Return, .Return_Void:
		return false
	}
	return false
}

write_operand :: proc(builder: ^strings.Builder, instructions: []ir.Instruction, value_id: ir.Instruction_Id, expected: types.Type) {
	if !valid_value(instructions, value_id, expected) {
		fmt.sbprintf(builder, "%d", sentinel(expected))
		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: source.Diagnostic_Id, span: source.Span, fallback: string) -> int {
	if _, ok := source.diagnostic_index(diagnostic, len(emitter.diagnostics.items)); ok {
		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_recovery_value :: proc(emitter: ^Emitter, instruction_id: int, instruction: ir.Instruction, fallback: string) {
	message := diagnostic_message(emitter, instruction.diagnostic, instruction.span, fallback)
	emit_trap_call(emitter, message)
	if types.is_concrete_integer(instruction.type) {
		fmt.sbprintf(
			&emitter.builder,
			"  %%v%d = add %s 0, %d\n",
			instruction_id,
			llvm_type(instruction.type),
			sentinel(instruction.type),
		)
	}
}

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

emit_instruction_stream :: proc(
	emitter: ^Emitter,
	instructions: []ir.Instruction,
	function: ir.Function,
	global_initializer := false,
) -> ir.Instruction_Id {
	return_value := ir.INVALID_INSTRUCTION
	after_return := false
	for instruction, instruction_index in instructions {
		instruction_id := ir.instruction_id(instruction_index)
		if after_return {
			fmt.sbprintf(&emitter.builder, "recover_after_return_%d:\n", instruction_index)
			after_return = false
		}
		switch instruction.op {
		case .Param, .Const:
		case .Load_Global:
			global_id := ir.as_global(instruction.target)
			if global_id == ir.INVALID_GLOBAL || int(global_id) >= len(emitter.module.globals) {
				emit_recovery_value(emitter, instruction_index, instruction, "invalid global reference")
				continue
			}
			global := emitter.module.globals[global_id]
			if !types.equal(instruction.type, global.type) {
				emit_recovery_value(emitter, instruction_index, instruction, "invalid global reference type")
				continue
			}
			if global.is_static {
				fmt.sbprintf(
					&emitter.builder,
					"  %%v%d = load %s, ptr @bro.g.%d\n",
					instruction_index,
					llvm_type(global.type),
					global_id,
				)
			} else {
				fmt.sbprintf(
					&emitter.builder,
					"  %%v%d = call %s @bro.get.%d()\n",
					instruction_index,
					llvm_type(global.type),
					global_id,
				)
			}
		case .Alloca:
			if !types.is_concrete_integer(instruction.type) {
				emit_recovery_value(emitter, instruction_index, instruction, "invalid allocation type")
				continue
			}
			fmt.sbprintf(&emitter.builder, "  %%v%d = alloca %s\n", instruction_index, llvm_type(instruction.type))
		case .Load:
			if !valid_instruction(instructions, instruction.a) ||
			   instructions[instruction.a].op != .Alloca ||
			   !types.equal(instructions[instruction.a].type, instruction.type) {
				emit_recovery_value(emitter, instruction_index, instruction, "invalid load slot")
				continue
			}
			fmt.sbprintf(
				&emitter.builder,
				"  %%v%d = load %s, ptr %%v%d\n",
				instruction_index,
				llvm_type(instruction.type),
				instruction.a,
			)
		case .Store:
			if !valid_instruction(instructions, instruction.a) ||
			   instructions[instruction.a].op != .Alloca ||
			   !types.equal(instructions[instruction.a].type, instruction.type) ||
			   !valid_value(instructions, instruction.b, instruction.type) {
				emit_recovery_value(emitter, instruction_index, instruction, "invalid store operand")
				continue
			}
			fmt.sbprintf(&emitter.builder, "  store %s ", llvm_type(instruction.type))
			write_operand(&emitter.builder, instructions, instruction.b, instruction.type)
			fmt.sbprintf(&emitter.builder, ", ptr %%v%d\n", instruction.a)
		case .Widen:
			if !valid_instruction(instructions, instruction.a) ||
			   !types.is_concrete_integer(instructions[instruction.a].type) ||
			   !types.is_concrete_integer(instruction.type) ||
			   instructions[instruction.a].type.bits >= instruction.type.bits {
				emit_recovery_value(emitter, instruction_index, instruction, "invalid widening operand")
				continue
			}
			from_type := instructions[instruction.a].type
			fmt.sbprintf(&emitter.builder, "  %%v%d = sext %s ", instruction_index, llvm_type(from_type))
			write_operand(&emitter.builder, instructions, instruction.a, from_type)
			fmt.sbprintf(&emitter.builder, " to %s\n", llvm_type(instruction.type))
		case .Neg_Checked:
			if !valid_value(instructions, instruction.a, instruction.type) {
				emit_recovery_value(emitter, instruction_index, instruction, "invalid negation operand")
				continue
			}
			type_name := llvm_type(instruction.type)
			fmt.sbprintf(&emitter.builder, "  %%pair%d = call ", instruction_index)
			strings.write_string(&emitter.builder, "{ ")
			fmt.sbprintf(&emitter.builder, "%s, i1 } @llvm.ssub.with.overflow.%s(%s 0, %s ", type_name, type_name, type_name, type_name)
			write_operand(&emitter.builder, instructions, instruction.a, instruction.type)
			fmt.sbprintf(&emitter.builder, ")\n")
			fmt.sbprintf(&emitter.builder, "  %%v%d = extractvalue ", instruction_index)
			strings.write_string(&emitter.builder, "{ ")
			fmt.sbprintf(&emitter.builder, "%s, i1 } %%pair%d, 0\n", type_name, instruction_index)
			fmt.sbprintf(&emitter.builder, "  %%overflow%d = extractvalue ", instruction_index)
			strings.write_string(&emitter.builder, "{ ")
			fmt.sbprintf(&emitter.builder, "%s, i1 } %%pair%d, 1\n", type_name, instruction_index)
			fmt.sbprintf(
				&emitter.builder,
				"  br i1 %%overflow%d, label %%overflow_trap%d, label %%overflow_continue%d\n",
				instruction_index,
				instruction_index,
				instruction_index,
			)
			fmt.sbprintf(&emitter.builder, "overflow_trap%d:\n", instruction_index)
			message := diagnostic_message(emitter, source.INVALID_DIAGNOSTIC, instruction.span, "signed integer negation overflow")
			emit_trap_call(emitter, message)
			fmt.sbprintf(&emitter.builder, "  unreachable\noverflow_continue%d:\n", instruction_index)
		case .Add_Checked:
			if !valid_value(instructions, instruction.a, instruction.type) ||
			   !valid_value(instructions, instruction.b, instruction.type) {
				emit_recovery_value(emitter, instruction_index, instruction, "invalid addition operand")
				continue
			}
			type_name := llvm_type(instruction.type)
			fmt.sbprintf(&emitter.builder, "  %%pair%d = call ", instruction_index)
			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, instruction.type)
			fmt.sbprintf(&emitter.builder, ", %s ", type_name)
			write_operand(&emitter.builder, instructions, instruction.b, instruction.type)
			fmt.sbprintf(&emitter.builder, ")\n")
			fmt.sbprintf(&emitter.builder, "  %%v%d = extractvalue ", instruction_index)
			strings.write_string(&emitter.builder, "{ ")
			fmt.sbprintf(&emitter.builder, "%s, i1 } %%pair%d, 0\n", type_name, instruction_index)
			fmt.sbprintf(&emitter.builder, "  %%overflow%d = extractvalue ", instruction_index)
			strings.write_string(&emitter.builder, "{ ")
			fmt.sbprintf(&emitter.builder, "%s, i1 } %%pair%d, 1\n", type_name, instruction_index)
			fmt.sbprintf(
				&emitter.builder,
				"  br i1 %%overflow%d, label %%overflow_trap%d, label %%overflow_continue%d\n",
				instruction_index,
				instruction_index,
				instruction_index,
			)
			fmt.sbprintf(&emitter.builder, "overflow_trap%d:\n", instruction_index)
			message := diagnostic_message(emitter, source.INVALID_DIAGNOSTIC, instruction.span, "signed integer addition overflow")
			emit_trap_call(emitter, message)
			fmt.sbprintf(&emitter.builder, "  unreachable\noverflow_continue%d:\n", instruction_index)
		case .Call:
			function_id := ir.as_function(instruction.target)
			if function_id == ir.INVALID_FUNCTION || int(function_id) >= len(emitter.module.functions) {
				emit_recovery_value(emitter, instruction_index, instruction, "invalid function specialization")
				continue
			}
			target := emitter.module.functions[function_id]
			valid_args := len(instruction.args) == len(target.param_types)
			if valid_args {
				for arg, index in instruction.args {
					if !valid_value(instructions, arg, target.param_types[index]) {
						valid_args = false
						break
					}
				}
			}
			target_result := target.result
			if target.is_main {
				target_result = types.I32
			}
			if !valid_args || !types.equal(instruction.type, target_result) {
				emit_recovery_value(emitter, instruction_index, instruction, "invalid function call operands")
				continue
			}
			if instruction.type.kind != .Void {
				fmt.sbprintf(&emitter.builder, "  %%v%d = ", instruction_index)
			} else {
				strings.write_string(&emitter.builder, "  ")
			}
			strings.write_string(&emitter.builder, "call ")
			if target.calling_convention == .Brolang {
				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, target.param_types)
			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, function.result)
			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, global.type)
		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, global.type)
		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 {
		if function.implementation == .Declaration {
			strings.write_string(&emitter.builder, "declare ")
		} else {
			strings.write_string(&emitter.builder, "define ")
			if function.linkage == .Internal {
				strings.write_string(&emitter.builder, "internal ")
			}
		}
		if function.calling_convention == .Brolang {
			strings.write_string(&emitter.builder, "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, ", ")
			}
			if function.implementation == .Declaration {
				fmt.sbprintf(&emitter.builder, "%s", llvm_type(param_type))
			} else {
				fmt.sbprintf(&emitter.builder, "%s %%v%d", llvm_type(param_type), index)
			}
		}
		if function.implementation == .Declaration {
			strings.write_string(&emitter.builder, ")\n\n")
			continue
		}
		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, "declare { i")
		fmt.sbprintf(&emitter.builder, "%d", bits)
		strings.write_string(&emitter.builder, ", i1 } @llvm.ssub.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) noreturn {\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)
}