// 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 macho_mh_magic_64 = u32(0xfeedfacf)
const macho_cpu_type_x86_64 = 0x01000007
const macho_cpu_subtype_x86_64_all = 3
const macho_mh_object = 1
const macho_lc_segment_64 = 0x19
const macho_lc_symtab = 0x2

const x86_64_reloc_signed = 1
const x86_64_reloc_branch = 2
const x86_64_reloc_got_load = 3
const x86_64_reloc_got = 4

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

	relocs      []MachORelocationInfo
	symbols     []MachOSymbol
	str_table   []u8
	sym_by_name map[string]int
}

struct MachORelocationInfo {
	addr    int
	sym_idx int
	pcrel   bool
	length  int
	extern  bool
	type_   int
}

struct MachOSymbol {
mut:
	name     string
	type_    u8
	sect     u8
	desc     u16
	value    u64
	name_off int
}

pub fn MachOObject.new() &MachOObject {
	return &MachOObject{
		str_table: [u8(0)]
	}
}

pub fn (mut m MachOObject) add_symbol(name string, value u64, is_ext bool, sect u8) int {
	typ := if is_ext { u8(0x0f) } else { u8(0x0e) }
	if i := m.sym_by_name[name] {
		mut s := &m.symbols[i]
		s.type_ = typ
		s.sect = sect
		s.value = value
		return i
	}

	idx := m.symbols.len
	name_off := m.add_string(name)
	m.symbols << MachOSymbol{
		name:     name
		type_:    typ
		sect:     sect
		desc:     0
		value:    value
		name_off: name_off
	}
	m.sym_by_name[name] = idx
	return idx
}

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

	idx := m.symbols.len
	name_off := m.add_string(name)
	m.symbols << MachOSymbol{
		name:     name
		type_:    0x01 // N_UNDF | N_EXT
		sect:     0
		desc:     0
		value:    0
		name_off: name_off
	}
	m.sym_by_name[name] = idx
	return idx
}

pub fn (mut m MachOObject) add_reloc(addr int, sym_idx int, typ int, pcrel bool, length int) {
	m.relocs << MachORelocationInfo{
		addr:    addr
		sym_idx: sym_idx
		type_:   typ
		pcrel:   pcrel
		length:  length
		extern:  true
	}
}

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

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

	n_sects := 3
	header_size := 32
	seg_cmd_size := 72 + (80 * n_sects)
	symtab_cmd_size := 24
	load_cmds_size := seg_cmd_size + symtab_cmd_size

	text_off := align_int(header_size + load_cmds_size, 16)
	text_len := m.text_data.len
	rodata_off := align_int(text_off + text_len, 8)
	rodata_len := m.rodata.len
	data_off := align_int(rodata_off + rodata_len, 8)
	data_len := m.data_data.len

	text_addr := 0
	rodata_addr := align_int(text_addr + text_len, 8)
	data_addr := align_int(rodata_addr + rodata_len, 8)
	total_size := data_addr + data_len

	reloc_off := align_int(data_off + data_len, 8)
	sym_off := align_int(reloc_off + (m.relocs.len * 8), 8)
	sym_len := m.symbols.len * 16
	str_off := sym_off + sym_len
	str_size := m.str_table.len

	write_u32_le(mut buf, macho_mh_magic_64)
	write_u32_le(mut buf, u32(macho_cpu_type_x86_64))
	write_u32_le(mut buf, u32(macho_cpu_subtype_x86_64_all))
	write_u32_le(mut buf, macho_mh_object)
	write_u32_le(mut buf, 2)
	write_u32_le(mut buf, u32(load_cmds_size))
	write_u32_le(mut buf, 0)
	write_u32_le(mut buf, 0)

	write_u32_le(mut buf, macho_lc_segment_64)
	write_u32_le(mut buf, u32(seg_cmd_size))
	write_fixed_string(mut buf, '', 16)
	write_u64_le(mut buf, 0)
	write_u64_le(mut buf, u64(total_size))
	write_u64_le(mut buf, u64(text_off))
	write_u64_le(mut buf, u64(data_off + data_len - text_off))
	write_u32_le(mut buf, 7)
	write_u32_le(mut buf, 7)
	write_u32_le(mut buf, u32(n_sects))
	write_u32_le(mut buf, 0)

	write_fixed_string(mut buf, '__text', 16)
	write_fixed_string(mut buf, '__TEXT', 16)
	write_u64_le(mut buf, u64(text_addr))
	write_u64_le(mut buf, u64(text_len))
	write_u32_le(mut buf, u32(text_off))
	write_u32_le(mut buf, 4)
	write_u32_le(mut buf, u32(reloc_off))
	write_u32_le(mut buf, u32(m.relocs.len))
	write_u32_le(mut buf, u32(0x80000400))
	write_u32_le(mut buf, 0)
	write_u32_le(mut buf, 0)
	write_u32_le(mut buf, 0)

	write_fixed_string(mut buf, '__const', 16)
	write_fixed_string(mut buf, '__TEXT', 16)
	write_u64_le(mut buf, u64(rodata_addr))
	write_u64_le(mut buf, u64(rodata_len))
	write_u32_le(mut buf, u32(rodata_off))
	write_u32_le(mut buf, 3)
	write_u32_le(mut buf, 0)
	write_u32_le(mut buf, 0)
	write_u32_le(mut buf, 0)
	write_u32_le(mut buf, 0)
	write_u32_le(mut buf, 0)
	write_u32_le(mut buf, 0)

	write_fixed_string(mut buf, '__data', 16)
	write_fixed_string(mut buf, '__DATA', 16)
	write_u64_le(mut buf, u64(data_addr))
	write_u64_le(mut buf, u64(data_len))
	write_u32_le(mut buf, u32(data_off))
	write_u32_le(mut buf, 3)
	write_u32_le(mut buf, 0)
	write_u32_le(mut buf, 0)
	write_u32_le(mut buf, 0)
	write_u32_le(mut buf, 0)
	write_u32_le(mut buf, 0)
	write_u32_le(mut buf, 0)

	write_u32_le(mut buf, macho_lc_symtab)
	write_u32_le(mut buf, u32(symtab_cmd_size))
	write_u32_le(mut buf, u32(sym_off))
	write_u32_le(mut buf, u32(m.symbols.len))
	write_u32_le(mut buf, u32(str_off))
	write_u32_le(mut buf, u32(str_size))

	for buf.len < text_off {
		buf << 0
	}
	buf << m.text_data
	for buf.len < rodata_off {
		buf << 0
	}
	buf << m.rodata
	for buf.len < data_off {
		buf << 0
	}
	buf << m.data_data

	for buf.len < reloc_off {
		buf << 0
	}
	for r in m.relocs {
		write_u32_le(mut buf, u32(r.addr))
		mut info := u32(r.sym_idx)
		if r.pcrel {
			info |= 1 << 24
		}
		info |= u32(r.length) << 25
		if r.extern {
			info |= 1 << 27
		}
		info |= u32(r.type_) << 28
		write_u32_le(mut buf, info)
	}

	for buf.len < sym_off {
		buf << 0
	}
	for s in m.symbols {
		write_u32_le(mut buf, u32(s.name_off))
		buf << s.type_
		buf << s.sect
		write_u16_le(mut buf, s.desc)
		write_u64_le(mut buf, m.symbol_value(s, rodata_addr, data_addr))
	}

	buf << m.str_table

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

fn (m MachOObject) symbol_value(s MachOSymbol, rodata_addr int, data_addr int) u64 {
	return match s.sect {
		1 { s.value }
		2 { u64(rodata_addr) + s.value }
		3 { u64(data_addr) + s.value }
		else { s.value }
	}
}

fn align_int(value int, alignment int) int {
	if alignment <= 1 || value % alignment == 0 {
		return value
	}
	return value + (alignment - (value % alignment))
}

fn write_fixed_string(mut b []u8, s string, len int) {
	bytes := s.bytes()
	for i in 0 .. len {
		if i < bytes.len {
			b << bytes[i]
		} else {
			b << 0
		}
	}
}