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

// ELF64 Constants
const ei_mag0 = 0x7f
const ei_mag1 = 0x45 // 'E'
const ei_mag2 = 0x4c // 'L'
const ei_mag3 = 0x46 // 'F'

const elfclass64 = 2
const elfdata2lsb = 1
const ev_current = 1
const et_rel = 1 // Relocatable file
const et_exec = 2 // Executable file
const em_x86_64 = 62

const pt_load = 1

const pf_x = 0x1
const pf_w = 0x2
const pf_r = 0x4

const sht_progbits = 1
const sht_symtab = 2
const sht_strtab = 3
const sht_rela = 4
const sht_nobits = 8

const shf_write = 0x1
const shf_alloc = 0x2
const shf_execinstr = 0x4

// Relocation Types for x86_64
const r_x86_64_64 = 1
const r_x86_64_pc32 = 2
const r_x86_64_plt32 = 4

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

	// Symbol Table
	symbols   []ElfSymbol
	str_table []u8

	// Relocations
	text_relocs []ElfRela

	// Section Names
	shstr_table []u8
}

struct ElfSymbol {
pub mut:
	name     string
	name_idx int
	info     u8
	shndx    u16
	value    u64
	size     u64
}

struct ElfRela {
	offset u64
	info   u64
	addend i64
}

pub fn ElfObject.new() &ElfObject {
	mut e := &ElfObject{
		str_table:   [u8(0)] // Starts with null byte
		shstr_table: [u8(0)]
	}
	// Add null symbol
	e.symbols << ElfSymbol{
		name:  ''
		shndx: 0
	}
	return e
}

pub fn (mut e ElfObject) add_symbol(name string, value u64, is_func bool, shndx u16) int {
	// STB_GLOBAL (1) << 4 | STT_FUNC (2) or STT_OBJECT (1) or STT_NOTYPE (0)
	type_ := if is_func { u8(2) } else { u8(1) } // Func : Object
	bind := u8(1) // Global
	info := (bind << 4) | type_

	for i := 1; i < e.symbols.len; i++ {
		if e.symbols[i].name != name {
			continue
		}
		mut s := &e.symbols[i]
		s.info = info
		s.shndx = shndx
		s.value = value
		return i
	}

	idx := e.symbols.len
	name_off := e.add_string(name)

	e.symbols << ElfSymbol{
		name:     name
		name_idx: name_off
		info:     info
		shndx:    shndx
		value:    value
	}
	return idx
}

pub fn (mut e ElfObject) add_undefined(name string) int {
	for i := 1; i < e.symbols.len; i++ {
		if e.symbols[i].name == name {
			return i
		}
	}

	idx := e.symbols.len
	name_off := e.add_string(name)
	// Bind Global (1), Type NoType (0)
	info := u8(0x10)

	e.symbols << ElfSymbol{
		name:     name
		name_idx: name_off
		info:     info
		shndx:    0 // Undefined
		value:    0
	}
	return idx
}

pub fn (mut e ElfObject) add_text_reloc(offset u64, sym_idx int, type_ int, addend i64) {
	// info = (sym_idx << 32) | type
	info := (u64(sym_idx) << 32) | u64(type_)
	e.text_relocs << ElfRela{
		offset: offset
		info:   info
		addend: addend
	}
}

fn (mut e ElfObject) add_string(s string) int {
	off := e.str_table.len
	e.str_table << s.bytes()
	e.str_table << 0
	return off
}

fn (mut e ElfObject) add_sh_string(s string) int {
	off := e.shstr_table.len
	e.shstr_table << s.bytes()
	e.shstr_table << 0
	return off
}

pub fn (mut e ElfObject) write(path string) {
	mut buf := []u8{}

	// --- 1. Prepare Sections ---
	// Indices:
	// 0: Null
	// 1: .text
	// 2: .data
	// 3: .rodata (strings)
	// 4: .symtab
	// 5: .strtab
	// 6: .rela.text
	// 7: .shstrtab

	// Prepare section names in .shstrtab
	off_text_name := e.add_sh_string('.text')
	off_data_name := e.add_sh_string('.data')
	off_rodata_name := e.add_sh_string('.rodata')
	off_symtab_name := e.add_sh_string('.symtab')
	off_strtab_name := e.add_sh_string('.strtab')
	off_rela_text_name := e.add_sh_string('.rela.text')
	off_shstrtab_name := e.add_sh_string('.shstrtab')

	// Calculate Offsets
	ehdr_size := 64
	shdr_entry_size := 64
	num_sections := 8

	mut current_offset := ehdr_size

	// .text (Align 16)
	if current_offset % 16 != 0 {
		current_offset += (16 - (current_offset % 16))
	}
	off_text := current_offset
	current_offset += e.text_data.len

	// .data (Align 8)
	if current_offset % 8 != 0 {
		current_offset += (8 - (current_offset % 8))
	}
	off_data := current_offset
	current_offset += e.data_data.len

	// .rodata (Align 4)
	if current_offset % 4 != 0 {
		current_offset += (4 - (current_offset % 4))
	}
	off_rodata := current_offset
	current_offset += e.rodata.len

	// .symtab (Align 8)
	if current_offset % 8 != 0 {
		current_offset += (8 - (current_offset % 8))
	}
	off_symtab := current_offset
	size_symtab := e.symbols.len * 24
	current_offset += size_symtab

	// .strtab (Align 1)
	off_strtab := current_offset
	current_offset += e.str_table.len

	// .rela.text (Align 8)
	if current_offset % 8 != 0 {
		current_offset += (8 - (current_offset % 8))
	}
	off_rela_text := current_offset
	size_rela_text := e.text_relocs.len * 24
	current_offset += size_rela_text

	// .shstrtab (Align 1)
	off_shstrtab := current_offset
	current_offset += e.shstr_table.len

	// Section Headers (Align 8)
	if current_offset % 8 != 0 {
		current_offset += (8 - (current_offset % 8))
	}
	off_shdrs := current_offset

	// --- 2. Write ELF Header ---
	buf << u8(ei_mag0)
	buf << u8(ei_mag1)
	buf << u8(ei_mag2)
	buf << u8(ei_mag3)
	buf << u8(elfclass64)
	buf << u8(elfdata2lsb)
	buf << u8(ev_current)
	buf << u8(0) // ABI System V
	buf << u8(0) // ABI Version
	for _ in 0 .. 7 {
		buf << 0
	} // Pad

	write_u16_le(mut buf, et_rel)
	write_u16_le(mut buf, em_x86_64)
	write_u32_le(mut buf, ev_current)
	write_u64_le(mut buf, 0) // Entry
	write_u64_le(mut buf, 0) // Phdr off
	write_u64_le(mut buf, u64(off_shdrs))
	write_u32_le(mut buf, 0) // Flags
	write_u16_le(mut buf, u16(ehdr_size))
	write_u16_le(mut buf, 0) // Phdr entry size
	write_u16_le(mut buf, 0) // Phdr num
	write_u16_le(mut buf, u16(shdr_entry_size))
	write_u16_le(mut buf, u16(num_sections))
	write_u16_le(mut buf, 7) // Shstrndx

	// --- 3. Write Data with Padding ---

	// Pad to Text
	for buf.len < off_text {
		buf << 0
	}
	buf << e.text_data

	// Pad to Data
	for buf.len < off_data {
		buf << 0
	}
	buf << e.data_data

	// Pad to Rodata
	for buf.len < off_rodata {
		buf << 0
	}
	buf << e.rodata

	// Pad to Symtab
	for buf.len < off_symtab {
		buf << 0
	}
	for sym in e.symbols {
		write_u32_le(mut buf, u32(sym.name_idx))
		buf << sym.info
		buf << 0 // other
		write_u16_le(mut buf, sym.shndx)
		write_u64_le(mut buf, sym.value)
		write_u64_le(mut buf, sym.size)
	}

	// Strtab
	for buf.len < off_strtab {
		buf << 0
	}
	buf << e.str_table

	// Rela Text
	for buf.len < off_rela_text {
		buf << 0
	}
	for r in e.text_relocs {
		write_u64_le(mut buf, r.offset)
		write_u64_le(mut buf, r.info)
		write_u64_le(mut buf, u64(r.addend))
	}

	// Shstrtab
	for buf.len < off_shstrtab {
		buf << 0
	}
	buf << e.shstr_table

	// --- 4. Write Section Headers ---
	for buf.len < off_shdrs {
		buf << 0
	}

	// 0: Null
	write_shdr(mut buf, 0, 0, 0, 0, 0, 0, 0, 0, 0)

	// 1: .text
	write_shdr(mut buf, u32(off_text_name), sht_progbits, shf_alloc | shf_execinstr, u64(off_text),
		u64(e.text_data.len), 0, 0, 16, 0)

	// 2: .data
	write_shdr(mut buf, u32(off_data_name), sht_progbits, shf_alloc | shf_write, u64(off_data),
		u64(e.data_data.len), 0, 0, 8, 0)

	// 3: .rodata
	write_shdr(mut buf, u32(off_rodata_name), sht_progbits, shf_alloc, u64(off_rodata),
		u64(e.rodata.len), 0, 0, 4, 0)

	// 4: .symtab (EntSize = 24)
	mut first_global := 1
	for i, s in e.symbols {
		if (s.info >> 4) == 1 { // STB_GLOBAL
			first_global = i
			break
		}
	}
	write_shdr(mut buf, u32(off_symtab_name), sht_symtab, 0, u64(off_symtab), u64(size_symtab), 5,
		u32(first_global), 8, 24)

	// 5: .strtab
	write_shdr(mut buf, u32(off_strtab_name), sht_strtab, 0, u64(off_strtab), u64(e.str_table.len),
		0, 0, 1, 0)

	// 6: .rela.text (EntSize = 24)
	write_shdr(mut buf, u32(off_rela_text_name), sht_rela, 0, u64(off_rela_text),
		u64(size_rela_text), 4, 1, 8, 24)

	// 7: .shstrtab
	write_shdr(mut buf, u32(off_shstrtab_name), sht_strtab, 0, u64(off_shstrtab),
		u64(e.shstr_table.len), 0, 0, 1, 0)

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

fn write_shdr(mut b []u8, name u32, type_ u32, flags u64, off u64, size u64, link u32, info u32, align u64, entsize u64) {
	write_u32_le(mut b, name)
	write_u32_le(mut b, type_)
	write_u64_le(mut b, flags)
	write_u64_le(mut b, 0) // Addr
	write_u64_le(mut b, off)
	write_u64_le(mut b, size)
	write_u32_le(mut b, link)
	write_u32_le(mut b, info)
	write_u64_le(mut b, align)
	write_u64_le(mut b, entsize)
}

fn write_u32_le(mut b []u8, v u32) {
	b << u8(v)
	b << u8(v >> 8)
	b << u8(v >> 16)
	b << u8(v >> 24)
}

fn write_u64_le(mut b []u8, v u64) {
	b << u8(v)
	b << u8(v >> 8)
	b << u8(v >> 16)
	b << u8(v >> 24)
	b << u8(v >> 32)
	b << u8(v >> 40)
	b << u8(v >> 48)
	b << u8(v >> 56)
}

fn write_u16_le(mut b []u8, v u16) {
	b << u8(v)
	b << u8(v >> 8)
}