// rasm-core/src/codegen/aot/mach_o.rs
//
//! # Mach-O 64-bit Writer
//!
//! Собирает минимальный исполняемый Mach-O для macOS x86_64 напрямую
//! из [`CompiledModule`](crate::object::CompiledModule):
//! каждая непустая секция модуля становится Mach-O section'ом внутри
//! соответствующего `LC_SEGMENT_64` (__TEXT/__DATA). `Abs64`-релокации
//! (`.quad my_label`) патчатся здесь, с учётом реальной базы загрузки.
//!
//! ## Структура файла
//!
//! ```text
//! Offset    Size  Content
//! ────────  ───── ────────────────────────────────────────
//! 0x00      32    mach_header_64
//! 0x20      72    LC_SEGMENT_64 __PAGEZERO
//! 0x68      72+   LC_SEGMENT_64 __TEXT + section_64 entries
//!           72+   LC_SEGMENT_64 __DATA + section_64 entries (if data/bss)
//!           24    LC_MAIN
//!           24    LC_BUILD_VERSION
//!           ...   padding to 0x1000
//!           ...   __text data
//!           ...   __data data (if any)
//! ```
use std::io::{self, Write};

use crate::{
    codegen::reloc::RelocKind,
    object::{CompiledModule, SectionKind},
};

// =========================================================================
// Mach-O constants
// =========================================================================

const MH_MAGIC_64: u32 = 0xFEED_FACF;
const MH_EXECUTE: u32 = 2;
const CPU_TYPE_X86_64: u32 = 0x0100_0007;
const CPU_SUBTYPE_X86_64_ALL: u32 = 3;
const MH_NOUNDEFS: u32 = 0x0001;
const MH_PIE: u32 = 0x0020_0000;

const LC_SEGMENT_64: u32 = 0x19;
const LC_MAIN: u32 = 0x28;
const LC_BUILD_VERSION: u32 = 0x32;

const PLATFORM_MACOS: u32 = 1;
const TEXT_BASE: u64 = 0x1_0000_0000;
const PAGE_SIZE: u64 = 0x1000;

const S_ATTR_PURE_INSTRUCTIONS: u32 = 0x8000_0000;
const S_ATTR_SOME_INSTRUCTIONS: u32 = 0x0000_0400;
const S_ZEROFILL: u32 = 0x0001;

const VM_PROT_READ: u32 = 1;
const VM_PROT_WRITE: u32 = 2;
const VM_PROT_EXECUTE: u32 = 4;

// =========================================================================
// MachOWriter
// =========================================================================

pub struct MachOWriter<'a> {
    module: &'a CompiledModule,
}

impl<'a> MachOWriter<'a> {
    pub fn new(module: &'a CompiledModule) -> Self {
        Self { module }
    }

    pub fn write<W: Write>(&self, out: &mut W) -> io::Result<()> {
        let ctx = MachOCtx::new(self.module);
        ctx.emit(out, self.module)
    }
}

// =========================================================================
// Internal layout structures
// =========================================================================

struct MachSection {
    /// Index into `module.sections`.
    module_index: usize,
    sectname: [u8; 16],
    segname: [u8; 16],
    addr: u64,
    size: u64,
    offset: u32,
    align_pow2: u32,
    flags: u32,
    has_file_data: bool,
}

struct MachOSegment {
    segname: [u8; 16],
    vmaddr: u64,
    vmsize: u64,
    fileoff: u64,
    filesize: u64,
    initprot: u32,
    sections: Vec<MachSection>,
}

struct MachOCtx {
    segments: Vec<MachOSegment>,
    entry_fileoff: u64,
    ncmds: u32,
    sizeofcmds: u32,
    data_seg: Option<usize>,
}

impl MachOCtx {
    fn new(module: &CompiledModule) -> Self {
        // --- Collect sections by kind ---
        let mut text_sections: Vec<(usize, &crate::object::Section)> = Vec::new();
        let mut data_file_sections: Vec<(usize, &crate::object::Section)> = Vec::new();
        let mut data_bss_sections: Vec<(usize, &crate::object::Section)> = Vec::new();

        for (idx, sec) in module.sections.iter().enumerate() {
            match sec.kind {
                SectionKind::Text => text_sections.push((idx, sec)),
                SectionKind::Data | SectionKind::Custom => data_file_sections.push((idx, sec)),
                SectionKind::Bss => data_bss_sections.push((idx, sec)),
            }
        }

        // --- Compute header sizes ---
        let ncmds = if data_file_sections.is_empty() && data_bss_sections.is_empty() {
            4u32 // PAGEZERO, TEXT, MAIN, BUILD_VERSION
        } else {
            5u32 // + DATA
        };

        let text_nsects = text_sections.len() as u32;
        let data_nsects = (data_file_sections.len() + data_bss_sections.len()) as u32;
        let has_data = data_nsects > 0;

        // sizeofcmds:
        //   PAGEZERO: 72
        //   TEXT:     72 + text_nsects * 80
        //   DATA:     72 + data_nsects * 80 (if has_data)
        //   MAIN:     24
        //   BUILD_VERSION: 24
        let mut sizeofcmds = 72u32; // PAGEZERO
        sizeofcmds += 72 + text_nsects * 80; // TEXT
        if has_data {
            sizeofcmds += 72 + data_nsects * 80; // DATA
        }
        sizeofcmds += 24 + 24; // MAIN + BUILD_VERSION

        // Headers size = mach_header_64 (32) + sizeofcmds
        let headers_total = 32u64 + sizeofcmds as u64;
        let headers_aligned = align_up(headers_total, PAGE_SIZE);

        // --- Build __TEXT segment ---
        let text_vmaddr = TEXT_BASE;
        let text_fileoff = headers_aligned;
        let mut text_bytes_total = 0u64;

        let mut text_mach_sections = Vec::new();
        for (idx, sec) in &text_sections {
            let size = sec.virtual_size as u64;
            text_mach_sections.push(MachSection {
                module_index: *idx,
                sectname: fixed_name(if sec.name.as_str() == ".text" {
                    "__text"
                } else {
                    &sec.name
                }),
                segname: fixed_name("__TEXT"),
                addr: text_vmaddr + text_bytes_total,
                size,
                offset: (text_fileoff + text_bytes_total) as u32,
                align_pow2: 2, // 2^2 = 4 bytes
                flags: S_ATTR_PURE_INSTRUCTIONS | S_ATTR_SOME_INSTRUCTIONS,
                has_file_data: true,
            });
            text_bytes_total += size;
        }

        let text_vmsize = align_up(text_bytes_total, PAGE_SIZE);
        let text_seg = MachOSegment {
            segname: fixed_name("__TEXT"),
            vmaddr: text_vmaddr,
            vmsize: text_vmsize,
            fileoff: text_fileoff,
            filesize: text_bytes_total,
            initprot: VM_PROT_READ | VM_PROT_EXECUTE,
            sections: text_mach_sections,
        };

        // --- Build __DATA segment (if needed) ---
        let mut data_seg: Option<MachOSegment> = None;
        if has_data {
            let data_vmaddr = TEXT_BASE + text_vmsize;
            let data_fileoff = align_up(text_fileoff + text_bytes_total, PAGE_SIZE);

            let mut data_mach_sections = Vec::new();
            let mut data_file_bytes_total = 0u64;

            for (idx, sec) in &data_file_sections {
                let size = sec.virtual_size as u64;
                data_mach_sections.push(MachSection {
                    module_index: *idx,
                    sectname: fixed_name(match sec.kind {
                        SectionKind::Data => "__data",
                        SectionKind::Custom => &sec.name,
                        _ => "__data",
                    }),
                    segname: fixed_name("__DATA"),
                    addr: data_vmaddr + data_file_bytes_total,
                    size,
                    offset: (data_fileoff + data_file_bytes_total) as u32,
                    align_pow2: 3, // 2^3 = 8 bytes
                    flags: 0,      // S_REGULAR
                    has_file_data: true,
                });
                data_file_bytes_total += size;
            }

            for (idx, sec) in &data_bss_sections {
                let size = sec.virtual_size as u64;
                data_mach_sections.push(MachSection {
                    module_index: *idx,
                    sectname: fixed_name("__bss"),
                    segname: fixed_name("__DATA"),
                    addr: data_vmaddr + data_file_bytes_total,
                    size,
                    offset: 0, // no file data
                    align_pow2: 3,
                    flags: S_ZEROFILL,
                    has_file_data: false,
                });
            }

            let data_vmsize = align_up(data_file_bytes_total, PAGE_SIZE);
            data_seg = Some(MachOSegment {
                segname: fixed_name("__DATA"),
                vmaddr: data_vmaddr,
                vmsize: data_vmsize,
                fileoff: data_fileoff,
                filesize: data_file_bytes_total,
                initprot: VM_PROT_READ | VM_PROT_WRITE,
                sections: data_mach_sections,
            });
        }

        // --- Entry point offset ---
        let entry_fileoff = Self::compute_entry_fileoff(module, text_fileoff, &text_sections);

        let mut segments = Vec::with_capacity(2);
        segments.push(text_seg);
        if let Some(ds) = data_seg {
            segments.push(ds);
        }

        let data_seg_idx = if has_data { Some(1) } else { None };

        Self {
            segments,
            entry_fileoff,
            ncmds,
            sizeofcmds,
            data_seg: data_seg_idx,
        }
    }

    fn compute_entry_fileoff(
        module: &CompiledModule,
        text_fileoff: u64,
        text_sections: &[(usize, &crate::object::Section)],
    ) -> u64 {
        let entry = match &module.entry {
            Some(e) => e,
            None => return 0,
        };

        // Find which text section and accumulate its file offset
        let mut acc = text_fileoff;
        for (idx, sec) in text_sections {
            if *idx == entry.section.0 {
                return acc + entry.offset as u64;
            }
            acc += sec.bytes.len() as u64;
        }

        0
    }

    fn symbol_vaddr(&self, module: &CompiledModule, name: &str) -> Option<u64> {
        let symbol = module.symbol(name)?;
        let section_id = symbol.section?;

        for seg in &self.segments {
            for mach_sec in &seg.sections {
                if mach_sec.module_index == section_id.0 {
                    return Some(mach_sec.addr + symbol.offset as u64);
                }
            }
        }

        None
    }

    fn patch_relocations(&self, module: &CompiledModule) -> Vec<Vec<u8>> {
        let mut bytes: Vec<Vec<u8>> = module.sections.iter().map(|s| s.bytes.clone()).collect();

        for reloc in &module.relocations {
            if reloc.kind != RelocKind::Abs64 {
                continue;
            }

            let Some(target_vaddr) = self.symbol_vaddr(module, &reloc.target) else {
                continue;
            };

            let section_bytes = &mut bytes[reloc.section.0];
            let offset = reloc.offset;
            if offset + 8 <= section_bytes.len() {
                section_bytes[offset..offset + 8].copy_from_slice(&target_vaddr.to_le_bytes());
            }
        }

        bytes
    }

    fn emit<W: Write>(&self, out: &mut W, module: &CompiledModule) -> io::Result<()> {
        let patched = self.patch_relocations(module);

        out.write_all(&self.header())?;
        out.write_all(&self.pagezero_segment())?;
        out.write_all(&self.segment_bytes(0))?;

        if let Some(ds) = self.data_seg {
            out.write_all(&self.segment_bytes(ds))?;
        }

        out.write_all(&self.main_command())?;
        out.write_all(&self.build_version())?;

        let mut written: u64 = 32
            + 72 // PAGEZERO
            + self.segment_bytes(0).len() as u64
            + self
                .data_seg
                .map(|ds| self.segment_bytes(ds).len() as u64)
                .unwrap_or(0)
            + 24 // LC_MAIN
            + 24; // LC_BUILD_VERSION

        for seg in &self.segments {
            pad_to(out, &mut written, seg.fileoff)?;

            for mach_sec in &seg.sections {
                if !mach_sec.has_file_data {
                    continue;
                }
                let idx = mach_sec.module_index;
                out.write_all(&patched[idx])?;
                written += mach_sec.size;
            }
        }

        Ok(())
    }

    // =====================================================================
    // Raw byte builders
    // =====================================================================

    fn header(&self) -> [u8; 32] {
        let mut h = [0u8; 32];
        put_u32(&mut h, 0, MH_MAGIC_64);
        put_u32(&mut h, 4, CPU_TYPE_X86_64);
        put_u32(&mut h, 8, CPU_SUBTYPE_X86_64_ALL);
        put_u32(&mut h, 12, MH_EXECUTE);
        put_u32(&mut h, 16, self.ncmds);
        put_u32(&mut h, 20, self.sizeofcmds);
        put_u32(&mut h, 24, MH_NOUNDEFS | MH_PIE);
        h
    }

    fn pagezero_segment(&self) -> [u8; 72] {
        let mut s = [0u8; 72];
        put_u32(&mut s, 0, LC_SEGMENT_64);
        put_u32(&mut s, 4, 72);
        s[8..24].copy_from_slice(&fixed_name("__PAGEZERO"));
        put_u64(&mut s, 24, 0);
        put_u64(&mut s, 32, TEXT_BASE); // cover 0..TEXT_BASE
        put_u64(&mut s, 40, 0);
        put_u64(&mut s, 48, 0);
        put_u32(&mut s, 56, 0);
        put_u32(&mut s, 60, 0);
        put_u32(&mut s, 64, 0);
        put_u32(&mut s, 68, 0);
        s
    }

    fn segment_bytes(&self, seg_idx: usize) -> Vec<u8> {
        let seg = &self.segments[seg_idx];
        let nsects = seg.sections.len() as u32;
        let total = 72 + nsects as usize * 80;
        let mut buf = vec![0u8; total];

        put_u32(&mut buf, 0, LC_SEGMENT_64);
        put_u32(&mut buf, 4, total as u32);
        buf[8..24].copy_from_slice(&seg.segname);
        put_u64(&mut buf, 24, seg.vmaddr);
        put_u64(&mut buf, 32, seg.vmsize);
        put_u64(&mut buf, 40, seg.fileoff);
        put_u64(&mut buf, 48, seg.filesize);
        put_u32(&mut buf, 56, seg.initprot); // maxprot
        put_u32(&mut buf, 60, seg.initprot); // initprot
        put_u32(&mut buf, 64, nsects);
        put_u32(&mut buf, 68, 0); // flags

        for (i, sec) in seg.sections.iter().enumerate() {
            let o = 72 + i * 80;
            buf[o..o + 16].copy_from_slice(&sec.sectname);
            buf[o + 16..o + 32].copy_from_slice(&sec.segname);
            put_u64(&mut buf, o + 32, sec.addr);
            put_u64(&mut buf, o + 40, sec.size);
            put_u32(&mut buf, o + 48, sec.offset);
            put_u32(&mut buf, o + 52, sec.align_pow2);
            put_u32(&mut buf, o + 56, 0); // reloff
            put_u32(&mut buf, o + 60, 0); // nreloc
            put_u32(&mut buf, o + 64, sec.flags);
            put_u32(&mut buf, o + 68, 0); // reserved1
            put_u32(&mut buf, o + 72, 0); // reserved2
            put_u32(&mut buf, o + 76, 0); // reserved3
        }

        buf
    }

    fn main_command(&self) -> [u8; 24] {
        let mut m = [0u8; 24];
        put_u32(&mut m, 0, LC_MAIN);
        put_u32(&mut m, 4, 24);
        put_u64(&mut m, 8, self.entry_fileoff);
        put_u64(&mut m, 16, 0); // stacksize = default
        m
    }

    fn build_version(&self) -> [u8; 24] {
        let mut b = [0u8; 24];
        put_u32(&mut b, 0, LC_BUILD_VERSION);
        put_u32(&mut b, 4, 24);
        put_u32(&mut b, 8, PLATFORM_MACOS);
        put_u32(&mut b, 12, 0x000E_0000); // macOS 14.0
        put_u32(&mut b, 16, 0x000E_0000); // SDK 14.0
        put_u32(&mut b, 20, 0); // ntools
        b
    }
}

// =========================================================================
// Helpers
// =========================================================================

fn fixed_name(name: &str) -> [u8; 16] {
    let mut buf = [0u8; 16];
    let bytes = name.as_bytes();
    let len = bytes.len().min(16);
    buf[..len].copy_from_slice(&bytes[..len]);
    buf
}

fn align_up(val: u64, align: u64) -> u64 {
    (val + align - 1) & !(align - 1)
}

fn pad_to<W: Write>(out: &mut W, written: &mut u64, target: u64) -> io::Result<()> {
    if *written >= target {
        return Ok(());
    }
    let zeros = vec![0u8; (target - *written) as usize];
    out.write_all(&zeros)?;
    *written = target;
    Ok(())
}

fn put_u32(buf: &mut [u8], offset: usize, val: u32) {
    buf[offset..offset + 4].copy_from_slice(&val.to_le_bytes());
}

fn put_u64(buf: &mut [u8], offset: usize, val: u64) {
    buf[offset..offset + 8].copy_from_slice(&val.to_le_bytes());
}
