// 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 os

const coff_image_file_machine_amd64 = u16(0x8664)
const coff_image_sym_class_external = u8(2)
const coff_image_sym_class_static = u8(3)
const coff_image_sym_dtype_function = u16(0x20)

const coff_image_rel_amd64_rel32 = u16(0x0004)

const coff_image_scn_cnt_code = u32(0x00000020)
const coff_image_scn_cnt_initialized_data = u32(0x00000040)
const coff_image_scn_align_8bytes = u32(0x00400000)
const coff_image_scn_align_16bytes = u32(0x00500000)
const coff_image_scn_mem_execute = u32(0x20000000)
const coff_image_scn_mem_read = u32(0x40000000)
const coff_image_scn_mem_write = u32(0x80000000)

pub struct CoffObject {
pub mut:
	text_data []u8
	data_data []u8
	rodata    []u8

	symbols     []CoffSymbol
	str_table   []u8
	sym_by_name map[string]int

	text_relocs []CoffRelocation
}

struct CoffSymbol {
mut:
	name          string
	name_off      int
	value         u32
	section       i16
	type_         u16
	storage_class u8
}

struct CoffRelocation {
	offset  u32
	sym_idx int
	type_   u16
}

struct CoffSection {
	name            string
	data            []u8
	relocs          []CoffRelocation
	characteristics u32
}

pub fn CoffObject.new() &CoffObject {
	return &CoffObject{}
}

pub fn (mut c CoffObject) add_symbol(name string, value u64, is_func bool, section u8) int {
	type_ := if is_func { coff_image_sym_dtype_function } else { u16(0) }
	storage_class := coff_symbol_storage_class(name, section)
	if i := c.sym_by_name[name] {
		mut s := &c.symbols[i]
		s.value = u32(value)
		s.section = i16(section)
		s.type_ = type_
		s.storage_class = storage_class
		return i
	}

	idx := c.symbols.len
	c.symbols << CoffSymbol{
		name:          name
		name_off:      c.add_string_if_needed(name)
		value:         u32(value)
		section:       i16(section)
		type_:         type_
		storage_class: storage_class
	}
	c.sym_by_name[name] = idx
	return idx
}

pub fn (mut c CoffObject) add_undefined(name string) int {
	if i := c.sym_by_name[name] {
		return i
	}

	idx := c.symbols.len
	c.symbols << CoffSymbol{
		name:          name
		name_off:      c.add_string_if_needed(name)
		value:         0
		section:       0
		type_:         0
		storage_class: coff_image_sym_class_external
	}
	c.sym_by_name[name] = idx
	return idx
}

fn coff_symbol_storage_class(name string, section u8) u8 {
	if section != 0 && name.starts_with('L_') {
		return coff_image_sym_class_static
	}
	return coff_image_sym_class_external
}

pub fn (mut c CoffObject) add_text_reloc(offset int, sym_idx int, type_ u16) {
	c.text_relocs << CoffRelocation{
		offset:  u32(offset)
		sym_idx: sym_idx
		type_:   type_
	}
}

fn (mut c CoffObject) add_string_if_needed(s string) int {
	if s.len <= 8 {
		return 0
	}
	off := 4 + c.str_table.len
	c.str_table << s.bytes()
	c.str_table << 0
	return off
}

pub fn (mut c CoffObject) write(path string) {
	sections := [
		CoffSection{
			name:            '.text'
			data:            c.text_data
			relocs:          c.text_relocs
			characteristics: coff_image_scn_cnt_code | coff_image_scn_mem_execute | coff_image_scn_mem_read | coff_image_scn_align_16bytes
		},
		CoffSection{
			name:            '.rdata'
			data:            c.rodata
			characteristics: coff_image_scn_cnt_initialized_data | coff_image_scn_mem_read | coff_image_scn_align_8bytes
		},
		CoffSection{
			name:            '.data'
			data:            c.data_data
			characteristics: coff_image_scn_cnt_initialized_data | coff_image_scn_mem_read | coff_image_scn_mem_write | coff_image_scn_align_8bytes
		},
	]

	header_size := 20
	section_header_size := 40
	reloc_entry_size := 10
	mut current_offset := header_size + (sections.len * section_header_size)

	mut raw_offsets := []int{len: sections.len}
	for i, section in sections {
		if section.data.len > 0 {
			current_offset = align_int(current_offset, 4)
			raw_offsets[i] = current_offset
			current_offset += section.data.len
		}
	}

	mut reloc_offsets := []int{len: sections.len}
	for i, section in sections {
		if section.relocs.len > 0 {
			current_offset = align_int(current_offset, 4)
			reloc_offsets[i] = current_offset
			current_offset += section.relocs.len * reloc_entry_size
		}
	}

	symbol_table_off := align_int(current_offset, 4)
	string_table_size := 4 + c.str_table.len

	mut buf := []u8{}
	write_u16_le(mut buf, coff_image_file_machine_amd64)
	write_u16_le(mut buf, u16(sections.len))
	write_u32_le(mut buf, 0)
	write_u32_le(mut buf, u32(symbol_table_off))
	write_u32_le(mut buf, u32(c.symbols.len))
	write_u16_le(mut buf, 0)
	write_u16_le(mut buf, 0)

	for i, section in sections {
		write_fixed_string(mut buf, section.name, 8)
		write_u32_le(mut buf, 0)
		write_u32_le(mut buf, 0)
		write_u32_le(mut buf, u32(section.data.len))
		write_u32_le(mut buf, u32(raw_offsets[i]))
		write_u32_le(mut buf, u32(reloc_offsets[i]))
		write_u32_le(mut buf, 0)
		nrelocs := coff_relocation_count_for_header(section) or { panic(err) }
		write_u16_le(mut buf, nrelocs)
		write_u16_le(mut buf, 0)
		write_u32_le(mut buf, section.characteristics)
	}

	for i, section in sections {
		if raw_offsets[i] == 0 {
			continue
		}
		for buf.len < raw_offsets[i] {
			buf << 0
		}
		buf << section.data
	}

	for i, section in sections {
		if reloc_offsets[i] == 0 {
			continue
		}
		for buf.len < reloc_offsets[i] {
			buf << 0
		}
		for reloc in section.relocs {
			write_u32_le(mut buf, reloc.offset)
			write_u32_le(mut buf, u32(reloc.sym_idx))
			write_u16_le(mut buf, reloc.type_)
		}
	}

	for buf.len < symbol_table_off {
		buf << 0
	}
	for sym in c.symbols {
		c.write_symbol_name(mut buf, sym)
		write_u32_le(mut buf, sym.value)
		write_u16_le(mut buf, u16(sym.section))
		write_u16_le(mut buf, sym.type_)
		buf << sym.storage_class
		buf << 0
	}

	write_u32_le(mut buf, u32(string_table_size))
	buf << c.str_table

	os.write_file_array(path, buf) or { panic(err) }
}

fn coff_relocation_count_for_header(section CoffSection) !u16 {
	if section.relocs.len > 0xffff {
		return error('COFF section ${section.name} has ${section.relocs.len} relocations; extended relocations are not supported')
	}
	return u16(section.relocs.len)
}

fn (c CoffObject) write_symbol_name(mut buf []u8, sym CoffSymbol) {
	if sym.name.len <= 8 {
		write_fixed_string(mut buf, sym.name, 8)
		return
	}
	write_u32_le(mut buf, 0)
	write_u32_le(mut buf, u32(sym.name_off))
}