// Copyright (c) 2026 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.

module x64

import encoding.binary

pub enum ObjectFormat {
	elf
	macho
	coff
}

enum ObjectSection {
	text
	data
	rodata
}

fn (format ObjectFormat) section_index(section ObjectSection) u8 {
	return match format {
		.elf {
			match section {
				.text { 1 }
				.data { 2 }
				.rodata { 3 }
			}
		}
		.macho {
			match section {
				.text { 1 }
				.rodata { 2 }
				.data { 3 }
			}
		}
		.coff {
			match section {
				.text { 1 }
				.rodata { 2 }
				.data { 3 }
			}
		}
	}
}

fn (format ObjectFormat) symbol_name(name string) string {
	if format == .macho && !name.starts_with('L_') {
		return '_' + name
	}
	return name
}

fn (mut g Gen) text_len() int {
	return match g.obj_format {
		.elf { g.elf.text_data.len }
		.macho { g.macho.text_data.len }
		.coff { g.coff.text_data.len }
	}
}

fn (mut g Gen) data_len() int {
	return match g.obj_format {
		.elf { g.elf.data_data.len }
		.macho { g.macho.data_data.len }
		.coff { g.coff.data_data.len }
	}
}

fn (mut g Gen) rodata_len() int {
	return match g.obj_format {
		.elf { g.elf.rodata.len }
		.macho { g.macho.rodata.len }
		.coff { g.coff.rodata.len }
	}
}

fn (mut g Gen) add_data_byte(b u8) {
	match g.obj_format {
		.elf { g.elf.data_data << b }
		.macho { g.macho.data_data << b }
		.coff { g.coff.data_data << b }
	}
}

fn (mut g Gen) add_data(bytes []u8) {
	match g.obj_format {
		.elf { g.elf.data_data << bytes }
		.macho { g.macho.data_data << bytes }
		.coff { g.coff.data_data << bytes }
	}
}

fn (mut g Gen) add_rodata_byte(b u8) {
	match g.obj_format {
		.elf { g.elf.rodata << b }
		.macho { g.macho.rodata << b }
		.coff { g.coff.rodata << b }
	}
}

fn (mut g Gen) add_rodata(bytes []u8) {
	match g.obj_format {
		.elf { g.elf.rodata << bytes }
		.macho { g.macho.rodata << bytes }
		.coff { g.coff.rodata << bytes }
	}
}

fn (mut g Gen) add_symbol(name string, value u64, is_func bool, section ObjectSection) int {
	sym_name := g.obj_format.symbol_name(name)
	sect := g.obj_format.section_index(section)
	return match g.obj_format {
		.elf { g.elf.add_symbol(sym_name, value, is_func, u16(sect)) }
		.macho { g.macho.add_symbol(sym_name, value, !sym_name.starts_with('L_'), sect) }
		.coff { g.coff.add_symbol(sym_name, value, is_func, sect) }
	}
}

fn (mut g Gen) add_undefined(name string) int {
	sym_name := g.obj_format.symbol_name(name)
	return match g.obj_format {
		.elf { g.elf.add_undefined(sym_name) }
		.macho { g.macho.add_undefined(sym_name) }
		.coff { g.coff.add_undefined(sym_name) }
	}
}

fn (mut g Gen) add_call_reloc(sym_idx int) {
	offset := g.text_len()
	match g.obj_format {
		.elf { g.elf.add_text_reloc(u64(offset), sym_idx, r_x86_64_plt32, -4) }
		.macho { g.macho.add_reloc(offset, sym_idx, x86_64_reloc_branch, true, 2) }
		.coff { g.coff.add_text_reloc(offset, sym_idx, coff_image_rel_amd64_rel32) }
	}
}

fn (mut g Gen) add_rip_reloc(sym_idx int) {
	offset := g.text_len()
	match g.obj_format {
		.elf { g.elf.add_text_reloc(u64(offset), sym_idx, r_x86_64_pc32, -4) }
		.macho { g.macho.add_reloc(offset, sym_idx, x86_64_reloc_signed, true, 2) }
		.coff { g.coff.add_text_reloc(offset, sym_idx, coff_image_rel_amd64_rel32) }
	}
}

fn (mut g Gen) add_macho_got_load_reloc(sym_idx int) {
	offset := g.text_len()
	if g.obj_format == .macho {
		g.macho.add_reloc(offset, sym_idx, x86_64_reloc_got_load, true, 2)
		return
	}
	x64_unsupported('Mach-O GOT_LOAD relocation requested for ${g.obj_format}')
}

fn (mut g Gen) emit(b u8) {
	match g.obj_format {
		.elf { g.elf.text_data << b }
		.macho { g.macho.text_data << b }
		.coff { g.coff.text_data << b }
	}
}

fn (mut g Gen) emit_u32(v u32) {
	g.emit(u8(v))
	g.emit(u8(v >> 8))
	g.emit(u8(v >> 16))
	g.emit(u8(v >> 24))
}

fn (mut g Gen) emit_u64(v u64) {
	g.emit_u32(u32(v))
	g.emit_u32(u32(v >> 32))
}

fn (mut g Gen) write_u32(off int, v u32) {
	match g.obj_format {
		.elf { binary.little_endian_put_u32(mut g.elf.text_data[off..off + 4], v) }
		.macho { binary.little_endian_put_u32(mut g.macho.text_data[off..off + 4], v) }
		.coff { binary.little_endian_put_u32(mut g.coff.text_data[off..off + 4], v) }
	}
}

pub fn (mut g Gen) write_file(path string) {
	if msg := g.unsupported_external_symbol_message() {
		x64_abort_with_diagnostic(msg)
	}
	match g.obj_format {
		.elf { g.elf.write(path) }
		.macho { g.macho.write(path) }
		.coff { g.coff.write(path) }
	}
}

fn (g Gen) unsupported_external_symbol_message() ?string {
	for raw_name in g.undefined_external_symbol_names() {
		if msg := unsupported_external_symbol_message_for_name(g.obj_format, raw_name,
			'needed while preparing native x64 output')
		{
			return msg
		}
	}
	return none
}

fn (g Gen) undefined_external_symbol_names() []string {
	mut names := []string{}
	match g.obj_format {
		.elf {
			for sym in g.elf.symbols {
				if sym.name != '' && sym.shndx == 0 {
					names << sym.name
				}
			}
		}
		.macho {
			for sym in g.macho.symbols {
				if sym.name != '' && sym.sect == 0 && sym.type_ == 0x01 {
					names << sym.name
				}
			}
		}
		.coff {
			for sym in g.coff.symbols {
				if sym.name != '' && sym.section == 0
					&& sym.storage_class == coff_image_sym_class_external {
					names << sym.name
				}
			}
		}
	}

	return names
}

pub fn (mut g Gen) link_executable(path string) ! {
	match g.obj_format {
		.coff {
			if g.abi != .windows {
				return error('x64 PE linking requires Windows ABI code generation')
			}
			mut linker := PeLinker.new(g.coff)
			linker.write(path)!
		}
		else {
			return error('x64 built-in executable linking is only implemented for COFF')
		}
	}
}