// 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 const macos_tiny_not_eligible_prefix = 'macOS tiny object is not eligible: '

const macos_tiny_builtin_string_plus_symbol = '_builtin__string__+'
const macos_tiny_builtin_i64_str_symbol = '_builtin__i64__str'
const macos_tiny_malloc_symbol = '_malloc'
const macos_tiny_exit_symbol = '_exit'

struct MachOTextRange {
    name  string
    start u64
    end   u64
}

struct MachODataRange {
    section ObjectSection
    start   u64
    end     u64
}

struct MachOTinyReachable {
mut:
    names     []string
    data      []MachODataRange
    undefined map[string]bool
}

struct MachOTinyObjectWriter {
    macho &MachOObject
}

struct MachOTinyRuntimeReloc {
    addr   int
    symbol string
}

pub fn (mut g Gen) write_macos_tiny_object(path string) ! {
    if g.obj_format != .macho {
        return error('macOS tiny object writing requires Mach-O object output')
    }
    if g.abi != .sysv {
        return error('macOS tiny object writing requires SysV x64 code generation')
    }
    mut writer := MachOTinyObjectWriter{
        macho: g.macho
    }
    writer.write(path)!
}

fn (mut w MachOTinyObjectWriter) write(path string) ! {
    mut reachable := w.collect_reachable()!
    mut out := MachOObject.new()
    mut symbol_indices := map[string]int{}
    mut func_offsets := map[string]u64{}
    mut runtime_symbols := []string{}
    mut runtime_relocs := []MachOTinyRuntimeReloc{}

    for name in reachable.names {
        range := w.text_range(name)!
        func_offsets[name] = u64(out.text_data.len)
        out.text_data << w.macho.text_data[int(range.start)..int(range.end)]
    }
    macho_tiny_add_runtime_helpers(mut out, mut reachable, mut func_offsets, mut runtime_symbols, mut
        runtime_relocs)

    mut data_offsets := map[string]u64{}
    w.copy_data_ranges(reachable.data, .rodata, mut out.rodata, mut data_offsets)!
    w.copy_data_ranges(reachable.data, .data, mut out.data_data, mut data_offsets)!

    for name in reachable.names {
        symbol_indices[name] = out.add_symbol(name, func_offsets[name],
            macho_tiny_defined_symbol_is_external(name), ObjectFormat.macho.section_index(.text))
    }
    for name in runtime_symbols {
        symbol_indices[name] = out.add_symbol(name, func_offsets[name],
            macho_tiny_defined_symbol_is_external(name), ObjectFormat.macho.section_index(.text))
    }
    mut data_names := data_offsets.keys()
    data_names.sort()
    for name in data_names {
        orig := w.symbol_by_name(name)!
        section := w.symbol_object_section(orig)!
        symbol_indices[name] = out.add_symbol(name, data_offsets[name],
            macho_tiny_defined_symbol_is_external(name), ObjectFormat.macho.section_index(section))
    }
    mut undefined_names := reachable.undefined.keys()
    undefined_names.sort()
    for name in undefined_names {
        symbol_indices[name] = out.add_undefined(name)
    }
    for name in reachable.names {
        range := w.text_range(name)!
        new_base := func_offsets[name]
        for reloc in w.macho.relocs {
            if reloc.addr < int(range.start) || reloc.addr >= int(range.end) {
                continue
            }
            sym := w.reloc_symbol(reloc)!
            sym_idx := symbol_indices[sym.name] or {
                return macos_tiny_not_eligible('symbol `${sym.name}` is not resolved by the tiny object')
            }
            out_reloc_addr := int(new_base + u64(reloc.addr - int(range.start)))
            if macho_tiny_resolve_relocation_locally(sym.name, reloc) {
                macho_tiny_patch_rel32_local(mut out.text_data, out_reloc_addr,
                    int(func_offsets[sym.name]))
                continue
            }
            out.add_reloc(out_reloc_addr, sym_idx, reloc.type_, reloc.pcrel, reloc.length)
        }
    }
    for reloc in runtime_relocs {
        sym_idx := symbol_indices[reloc.symbol] or {
            return macos_tiny_not_eligible('runtime helper symbol `${reloc.symbol}` is not resolved by the tiny object')
        }
        out.add_reloc(reloc.addr, sym_idx, x86_64_reloc_branch, true, 2)
    }

    out.write(path)
}

fn (mut w MachOTinyObjectWriter) collect_reachable() !MachOTinyReachable {
    _ = w.text_range('_main') or { return macos_tiny_not_eligible('missing _main symbol') }
    mut selected := []string{}
    mut seen := map[string]bool{}
    mut queue := ['_main']
    mut selected_data := []MachODataRange{}
    mut undefined := map[string]bool{}
    for queue.len > 0 {
        name := queue[0]
        queue.delete(0)
        if seen[name] {
            continue
        }
        if macho_tiny_is_module_init_symbol(name) {
            return macos_tiny_not_eligible('module init symbol `${name}` is not supported')
        }
        range := w.text_range(name)!
        seen[name] = true
        selected << name
        for reloc in w.macho.relocs {
            if reloc.addr < int(range.start) || reloc.addr >= int(range.end) {
                continue
            }
            validate_macho_tiny_relocation(reloc)!
            sym := w.reloc_symbol(reloc)!
            if sym.sect == ObjectFormat.macho.section_index(.text) {
                if !seen[sym.name] {
                    queue << sym.name
                }
            } else if sym.sect == ObjectFormat.macho.section_index(.rodata) {
                w.add_data_range(mut selected_data, .rodata, sym.value)!
            } else if sym.sect == ObjectFormat.macho.section_index(.data) {
                if x64_main_argc_global_name(sym.name) || x64_main_argv_global_name(sym.name) {
                    return macos_tiny_not_eligible('arguments() requires hosted argc/argv; macOS tiny _main argv is not implemented yet')
                }
                w.add_data_range(mut selected_data, .data, sym.value)!
            } else if sym.sect == 0 && sym.type_ == 0x01 {
                undefined[sym.name] = true
            } else {
                return macos_tiny_not_eligible('relocation to symbol `${sym.name}` in Mach-O section ${sym.sect} is not supported')
            }
        }
    }
    return MachOTinyReachable{
        names:     selected
        data:      selected_data
        undefined: undefined
    }
}

fn validate_macho_tiny_relocation(reloc MachORelocationInfo) ! {
    if !reloc.pcrel || reloc.length != 2 {
        return macos_tiny_not_eligible('non-pcrel or non-32-bit relocation at __text+${reloc.addr} is not supported')
    }
    if reloc.type_ !in [x86_64_reloc_signed, x86_64_reloc_branch, x86_64_reloc_got_load,
        x86_64_reloc_got] {
        return macos_tiny_not_eligible('Mach-O relocation type ${reloc.type_} is not supported')
    }
}

fn (w MachOTinyObjectWriter) text_range(name string) !MachOTextRange {
    ranges := w.text_symbol_ranges()
    for range in ranges {
        if range.name == name {
            return range
        }
    }
    return error('missing text symbol `${name}`')
}

fn (w MachOTinyObjectWriter) text_symbol_ranges() []MachOTextRange {
    mut syms := []MachOSymbol{}
    for sym in w.macho.symbols {
        if sym.name != '' && sym.sect == ObjectFormat.macho.section_index(.text) {
            syms << sym
        }
    }
    for i := 1; i < syms.len; i++ {
        mut j := i
        for j > 0 && syms[j - 1].value > syms[j].value {
            syms[j - 1], syms[j] = syms[j], syms[j - 1]
            j--
        }
    }
    mut ranges := []MachOTextRange{}
    for i, sym in syms {
        mut end := u64(w.macho.text_data.len)
        for j := i + 1; j < syms.len; j++ {
            if syms[j].value > sym.value {
                end = syms[j].value
                break
            }
        }
        if end > sym.value {
            ranges << MachOTextRange{
                name:  sym.name
                start: sym.value
                end:   end
            }
        }
    }
    return ranges
}

fn (w MachOTinyObjectWriter) reloc_symbol(reloc MachORelocationInfo) !MachOSymbol {
    if reloc.sym_idx < 0 || reloc.sym_idx >= w.macho.symbols.len {
        return error('macOS tiny object relocation references invalid symbol index ${reloc.sym_idx}')
    }
    return w.macho.symbols[reloc.sym_idx]
}

fn (w MachOTinyObjectWriter) symbol_by_name(name string) !MachOSymbol {
    idx := w.macho.sym_by_name[name] or { return error('missing symbol `${name}`') }
    if idx < 0 || idx >= w.macho.symbols.len {
        return error('invalid symbol index ${idx} for `${name}`')
    }
    return w.macho.symbols[idx]
}

fn (w MachOTinyObjectWriter) symbol_object_section(sym MachOSymbol) !ObjectSection {
    if sym.sect == ObjectFormat.macho.section_index(.rodata) {
        return .rodata
    }
    if sym.sect == ObjectFormat.macho.section_index(.data) {
        return .data
    }
    if sym.sect == ObjectFormat.macho.section_index(.text) {
        return .text
    }
    return error('symbol `${sym.name}` is not in a known Mach-O section')
}

fn (mut w MachOTinyObjectWriter) add_data_range(mut ranges []MachODataRange, section ObjectSection, start u64) ! {
    end := w.data_symbol_range_end(section, start)!
    for range in ranges {
        if range.section == section && range.start == start && range.end == end {
            return
        }
    }
    ranges << MachODataRange{
        section: section
        start:   start
        end:     end
    }
}

fn (w MachOTinyObjectWriter) data_symbol_range_end(section ObjectSection, start u64) !u64 {
    section_idx := ObjectFormat.macho.section_index(section)
    section_len := match section {
        .rodata { u64(w.macho.rodata.len) }
        .data { u64(w.macho.data_data.len) }
        else { u64(0) }
    }

    mut end := section_len
    for sym in w.macho.symbols {
        if sym.sect == section_idx && sym.value > start && sym.value < end {
            end = sym.value
        }
    }
    if end < start {
        return error('invalid Mach-O data symbol range')
    }
    return end
}

fn (w MachOTinyObjectWriter) copy_data_ranges(ranges []MachODataRange, section ObjectSection, mut out []u8, mut offsets map[string]u64) ! {
    for dr in ranges {
        if dr.section != section {
            continue
        }
        align := int(w.data_section_alignment(section))
        for align > 1 && out.len % align != int(dr.start % u64(align)) {
            out << u8(0)
        }
        off := u64(out.len)
        section_idx := ObjectFormat.macho.section_index(section)
        for sym in w.macho.symbols {
            if sym.sect == section_idx && sym.value == dr.start {
                offsets[sym.name] = off
            }
        }
        match section {
            .rodata {
                out << w.macho.rodata[int(dr.start)..int(dr.end)]
            }
            .data {
                out << w.macho.data_data[int(dr.start)..int(dr.end)]
            }
            else {
                return error('macOS tiny object cannot copy ${section} as data')
            }
        }
    }
}

fn (w MachOTinyObjectWriter) data_section_alignment(section ObjectSection) u64 {
    return match section {
        .data { u64(8) }
        .rodata { u64(4) }
        else { u64(1) }
    }
}

fn macho_tiny_defined_symbol_is_external(name string) bool {
    return name == '_main'
}

fn macho_tiny_resolve_relocation_locally(name string, reloc MachORelocationInfo) bool {
    return name in [macos_tiny_builtin_string_plus_symbol, macos_tiny_builtin_i64_str_symbol]
        && reloc.type_ == x86_64_reloc_branch && reloc.pcrel && reloc.length == 2
}

fn macho_tiny_add_runtime_helpers(mut out MachOObject, mut reachable MachOTinyReachable, mut func_offsets map[string]u64, mut runtime_symbols []string, mut runtime_relocs []MachOTinyRuntimeReloc) {
    if reachable.undefined[macos_tiny_builtin_i64_str_symbol] {
        reachable.undefined.delete(macos_tiny_builtin_i64_str_symbol)
        reachable.undefined[macos_tiny_malloc_symbol] = true
        reachable.undefined[macos_tiny_exit_symbol] = true
        func_offsets[macos_tiny_builtin_i64_str_symbol] = u64(out.text_data.len)
        runtime_symbols << macos_tiny_builtin_i64_str_symbol
        macho_tiny_emit_i64_str(mut out.text_data, mut runtime_relocs)
    }
    if reachable.undefined[macos_tiny_builtin_string_plus_symbol] {
        reachable.undefined.delete(macos_tiny_builtin_string_plus_symbol)
        reachable.undefined[macos_tiny_malloc_symbol] = true
        reachable.undefined[macos_tiny_exit_symbol] = true
        func_offsets[macos_tiny_builtin_string_plus_symbol] = u64(out.text_data.len)
        runtime_symbols << macos_tiny_builtin_string_plus_symbol
        macho_tiny_emit_string_plus(mut out.text_data, mut runtime_relocs)
    }
}

fn macho_tiny_emit_i64_str(mut text []u8, mut runtime_relocs []MachOTinyRuntimeReloc) {
    text << [
        u8(0x57), // push rdi, input value
        0xbf,
        0x20,
        0x00,
        0x00,
        0x00, // mov edi, 32
    ]
    macho_tiny_emit_call_reloc(mut text, mut runtime_relocs, macos_tiny_malloc_symbol)
    text << [
        u8(0x48),
        0x85,
        0xc0, // test rax, rax
    ]
    malloc_ok := macho_tiny_emit_jcc32(mut text, 0x85) // jne ok
    text << [
        u8(0xbf),
        0x01,
        0x00,
        0x00,
        0x00, // mov edi, 1
    ]
    macho_tiny_emit_call_reloc(mut text, mut runtime_relocs, macos_tiny_exit_symbol)
    text << [u8(0x0f), 0x0b] // ud2
    ok_start := text.len
    macho_tiny_patch_rel32_local(mut text, malloc_ok, ok_start)

    text << [
        u8(0x5f), // pop rdi
        0x4c,
        0x8d,
        0x40,
        0x1f, // lea r8, [rax+31]
        0x41,
        0xc6,
        0x00,
        0x00, // mov byte ptr [r8], 0
        0x48,
        0x89,
        0xf8, // mov rax, rdi
        0x45,
        0x31,
        0xc9, // xor r9d, r9d
        0x45,
        0x31,
        0xd2, // xor r10d, r10d
        0x48,
        0x85,
        0xc0, // test rax, rax
    ]
    non_negative := macho_tiny_emit_jcc32(mut text, 0x89) // jns non_negative
    text << [
        u8(0x41),
        0xb2,
        0x01, // mov r10b, 1
        0x48,
        0xf7,
        0xd8, // neg rax
    ]
    non_negative_target := text.len
    text << [
        u8(0x48),
        0x85,
        0xc0, // test rax, rax
    ]
    non_zero := macho_tiny_emit_jcc32(mut text, 0x85) // jne loop
    text << [
        u8(0x49),
        0xff,
        0xc8, // dec r8
        0x41,
        0xc6,
        0x00,
        0x30, // mov byte ptr [r8], '0'
        0x41,
        0xb9,
        0x01,
        0x00,
        0x00,
        0x00, // mov r9d, 1
    ]
    maybe_sign_jump := macho_tiny_emit_jmp32(mut text)
    loop_start := text.len
    text << [
        u8(0x31),
        0xd2, // xor edx, edx
        0xb9,
        0x0a,
        0x00,
        0x00,
        0x00, // mov ecx, 10
        0x48,
        0xf7,
        0xf1, // div rcx
        0x80,
        0xc2,
        0x30, // add dl, '0'
        0x49,
        0xff,
        0xc8, // dec r8
        0x41,
        0x88,
        0x10, // mov [r8], dl
        0x49,
        0xff,
        0xc1, // inc r9
        0x48,
        0x85,
        0xc0, // test rax, rax
    ]
    loop_more := macho_tiny_emit_jcc32(mut text, 0x85) // jne loop
    maybe_sign := text.len
    text << [
        u8(0x45),
        0x84,
        0xd2, // test r10b, r10b
    ]
    done_digits := macho_tiny_emit_jcc32(mut text, 0x84) // je done
    text << [
        u8(0x49),
        0xff,
        0xc8, // dec r8
        0x41,
        0xc6,
        0x00,
        0x2d, // mov byte ptr [r8], '-'
        0x49,
        0xff,
        0xc1, // inc r9
    ]
    done_digits_target := text.len
    text << [
        u8(0x4c),
        0x89,
        0xc0, // mov rax, r8
        0x4c,
        0x89,
        0xca, // mov rdx, r9
        0xc3, // ret
    ]
    macho_tiny_patch_rel32_local(mut text, non_negative, non_negative_target)
    macho_tiny_patch_rel32_local(mut text, non_zero, loop_start)
    macho_tiny_patch_rel32_local(mut text, maybe_sign_jump, maybe_sign)
    macho_tiny_patch_rel32_local(mut text, loop_more, loop_start)
    macho_tiny_patch_rel32_local(mut text, done_digits, done_digits_target)
}

fn macho_tiny_emit_string_plus(mut text []u8, mut runtime_relocs []MachOTinyRuntimeReloc) {
    text << [
        u8(0x57), // push rdi, first string ptr
        0x56, // push rsi, first string len
        0x52, // push rdx, second string ptr
        0x51, // push rcx, second string len
        0x48,
        0x83,
        0xec,
        0x08, // sub rsp, 8, align before libSystem calls
        0x44,
        0x8b,
        0x44,
        0x24,
        0x18, // mov r8d, [rsp+24]
        0x44,
        0x03,
        0x44,
        0x24,
        0x08, // add r8d, [rsp+8]
    ]
    len_overflow := macho_tiny_emit_jcc32(mut text, 0x82) // jc fail
    text << [
        u8(0x4d),
        0x89,
        0xc1, // mov r9, r8
        0x49,
        0x83,
        0xc1,
        0x01, // add r9, 1
    ]
    alloc_overflow := macho_tiny_emit_jcc32(mut text, 0x82) // jc fail
    text << [
        u8(0x4c),
        0x89,
        0xcf, // mov rdi, r9
    ]
    macho_tiny_emit_call_reloc(mut text, mut runtime_relocs, macos_tiny_malloc_symbol)
    text << [
        u8(0x48),
        0x85,
        0xc0, // test rax, rax
    ]
    malloc_ok := macho_tiny_emit_jcc32(mut text, 0x85) // jne ok
    fail_start := text.len
    text << [
        u8(0xbf),
        0x01,
        0x00,
        0x00,
        0x00, // mov edi, 1
    ]
    macho_tiny_emit_call_reloc(mut text, mut runtime_relocs, macos_tiny_exit_symbol)
    text << [u8(0x0f), 0x0b] // ud2
    ok_start := text.len
    macho_tiny_patch_rel32_local(mut text, len_overflow, fail_start)
    macho_tiny_patch_rel32_local(mut text, alloc_overflow, fail_start)
    macho_tiny_patch_rel32_local(mut text, malloc_ok, ok_start)

    text << [
        u8(0x49),
        0x89,
        0xc2, // mov r10, rax
        0x49,
        0x89,
        0xc0, // mov r8, rax
        0x48,
        0x8b,
        0x74,
        0x24,
        0x20, // mov rsi, [rsp+32]
        0x8b,
        0x4c,
        0x24,
        0x18, // mov ecx, [rsp+24]
        0x48,
        0x85,
        0xc9, // test rcx, rcx
    ]
    copy_a_done := macho_tiny_emit_jcc32(mut text, 0x84) // jz done
    copy_a_loop := text.len
    text << [
        u8(0x8a),
        0x16, // mov dl, [rsi]
        0x41,
        0x88,
        0x10, // mov [r8], dl
        0x48,
        0xff,
        0xc6, // inc rsi
        0x49,
        0xff,
        0xc0, // inc r8
        0x48,
        0xff,
        0xc9, // dec rcx
    ]
    copy_a_continue := macho_tiny_emit_jcc32(mut text, 0x85) // jne loop
    copy_a_done_target := text.len
    macho_tiny_patch_rel32_local(mut text, copy_a_done, copy_a_done_target)
    macho_tiny_patch_rel32_local(mut text, copy_a_continue, copy_a_loop)

    text << [
        u8(0x48),
        0x8b,
        0x74,
        0x24,
        0x10, // mov rsi, [rsp+16]
        0x8b,
        0x4c,
        0x24,
        0x08, // mov ecx, [rsp+8]
        0x48,
        0x85,
        0xc9, // test rcx, rcx
    ]
    copy_b_done := macho_tiny_emit_jcc32(mut text, 0x84) // jz done
    copy_b_loop := text.len
    text << [
        u8(0x8a),
        0x16, // mov dl, [rsi]
        0x41,
        0x88,
        0x10, // mov [r8], dl
        0x48,
        0xff,
        0xc6, // inc rsi
        0x49,
        0xff,
        0xc0, // inc r8
        0x48,
        0xff,
        0xc9, // dec rcx
    ]
    copy_b_continue := macho_tiny_emit_jcc32(mut text, 0x85) // jne loop
    copy_b_done_target := text.len
    macho_tiny_patch_rel32_local(mut text, copy_b_done, copy_b_done_target)
    macho_tiny_patch_rel32_local(mut text, copy_b_continue, copy_b_loop)

    text << [
        u8(0x41),
        0xc6,
        0x00,
        0x00, // mov byte ptr [r8], 0
        0x8b,
        0x44,
        0x24,
        0x18, // mov eax, [rsp+24]
        0x03,
        0x44,
        0x24,
        0x08, // add eax, [rsp+8]
        0x48,
        0x89,
        0xc2, // mov rdx, rax
        0x4c,
        0x89,
        0xd0, // mov rax, r10
        0x48,
        0x83,
        0xc4,
        0x28, // add rsp, 40
        0xc3, // ret
    ]
}

fn macho_tiny_emit_call_reloc(mut text []u8, mut runtime_relocs []MachOTinyRuntimeReloc, symbol string) {
    text << u8(0xe8)
    field_off := text.len
    text << [u8(0), 0, 0, 0]
    runtime_relocs << MachOTinyRuntimeReloc{
        addr:   field_off
        symbol: symbol
    }
}

fn macho_tiny_emit_jcc32(mut text []u8, opcode u8) int {
    text << [u8(0x0f), opcode]
    field_off := text.len
    text << [u8(0), 0, 0, 0]
    return field_off
}

fn macho_tiny_emit_jmp32(mut text []u8) int {
    text << u8(0xe9)
    field_off := text.len
    text << [u8(0), 0, 0, 0]
    return field_off
}

fn macho_tiny_patch_rel32_local(mut text []u8, field_off int, target_off int) {
    disp := i64(target_off) - i64(field_off + 4)
    binary.little_endian_put_u32(mut text[field_off..field_off + 4], u32(i32(disp)))
}

fn macho_tiny_is_module_init_symbol(name string) bool {
    return name.ends_with('__init') && name != '_main'
}

fn macos_tiny_not_eligible(message string) IError {
    return error(macos_tiny_not_eligible_prefix + message)
}