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

// import encoding.binary
import os

// Mach-O Constants for ARM64
const mh_magic_64 = u32(0xfeedfacf)
const cpu_type_arm64 = 0x0100000c
const cpu_subtype_arm64_all = 0

const lc_segment_64 = 0x19
const lc_symtab = 0x2

const arm64_reloc_branch26 = 2
const arm64_reloc_page21 = 3
const arm64_reloc_pageoff12 = 4
const arm64_reloc_got_load_page21 = 5
const arm64_reloc_got_load_pageoff12 = 6

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

    relocs      []RelocationInfo
    symbols     []Symbol
    str_table   []u8
    sym_by_name map[string]int // symbol name → index in symbols
}

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

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

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

pub fn (mut m MachOObject) add_symbol(name string, addr u64, is_ext bool, sect u8) int {
    typ := if is_ext { u8(0x0f) } else { u8(0x0e) } // N_SECT | N_EXT : N_SECT

    // Check if symbol already exists (e.g., was added as undefined earlier)
    if i := m.sym_by_name[name] {
        mut s := &m.symbols[i]
        s.type_ = typ
        s.sect = sect
        s.value = addr
        return i
    }

    // Add new symbol
    idx := m.symbols.len
    name_off := m.str_table.len
    unsafe { m.str_table.push_many(name.str, name.len) }
    m.str_table << 0

    m.symbols << Symbol{
        name:     name
        type_:    typ
        sect:     sect
        desc:     0
        value:    addr
        name_off: name_off
    }
    m.sym_by_name[name] = idx
    return idx
}

pub fn (mut m MachOObject) add_undefined(name string) int {
    // Check for any existing symbol with this name (defined or undefined)
    if i := m.sym_by_name[name] {
        return i
    }

    idx := m.symbols.len
    name_off := m.str_table.len
    unsafe { m.str_table.push_many(name.str, name.len) }
    m.str_table << 0

    m.symbols << Symbol{
        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) {
    m.relocs << RelocationInfo{
        addr:    addr
        sym_idx: sym_idx // 0-based symbol table index
        type_:   typ
        pcrel:   pcrel
        length:  2
        extern:  true
    }
}

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 := header_size + load_cmds_size
    text_len := m.text_data.len
    cstring_off := text_off + text_len
    cstring_len := m.str_data.len
    data_off := cstring_off + cstring_len
    data_len := m.data_data.len

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

    // 1. Header
    write_u32_le(mut buf, mh_magic_64)
    write_u32_le(mut buf, u32(cpu_type_arm64))
    write_u32_le(mut buf, u32(cpu_subtype_arm64_all))
    write_u32_le(mut buf, 1) // MH_OBJECT
    write_u32_le(mut buf, 2) // ncmds
    write_u32_le(mut buf, u32(load_cmds_size))
    write_u32_le(mut buf, 0)
    write_u32_le(mut buf, 0)

    // 2. LC_SEGMENT_64
    write_u32_le(mut buf, u32(lc_segment_64))
    write_u32_le(mut buf, u32(seg_cmd_size))
    for _ in 0 .. 16 {
        buf << 0
    }
    write_u64_le(mut buf, 0)
    write_u64_le(mut buf, u64(text_len + cstring_len + data_len))
    write_u64_le(mut buf, u64(text_off))
    write_u64_le(mut buf, u64(text_len + cstring_len + data_len))
    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)

    // Section 1: __text
    write_string_fixed(mut buf, '__text', 16)
    write_string_fixed(mut buf, '__TEXT', 16)
    write_u64_le(mut buf, 0) // Addr 0
    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)

    // Section 2: __cstring
    write_string_fixed(mut buf, '__cstring', 16)
    write_string_fixed(mut buf, '__TEXT', 16)
    write_u64_le(mut buf, u64(text_len)) // Addr = text_len
    write_u64_le(mut buf, u64(cstring_len))
    write_u32_le(mut buf, u32(cstring_off))
    write_u32_le(mut buf, 0)
    write_u32_le(mut buf, 0)
    write_u32_le(mut buf, 0)
    write_u32_le(mut buf, 2) // S_CSTRING_LITERALS
    write_u32_le(mut buf, 0)
    write_u32_le(mut buf, 0)
    write_u32_le(mut buf, 0)

    // Section 3: __data
    write_string_fixed(mut buf, '__data', 16)
    write_string_fixed(mut buf, '__DATA', 16)
    write_u64_le(mut buf, u64(text_len + cstring_len)) // Addr = text_len + cstring_len
    write_u64_le(mut buf, u64(data_len))
    write_u32_le(mut buf, u32(data_off))
    write_u32_le(mut buf, 3) // align 8
    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)

    // 3. LC_SYMTAB
    write_u32_le(mut buf, u32(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))

    buf << m.text_data
    buf << m.str_data
    buf << m.data_data

    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 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, s.value)
    }

    buf << m.str_table

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

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

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