// rasm-core/src/codegen/aot/elf.rs
//
//! # ELF64 Writer
//!
//! Собирает минимальный исполняемый ELF64-файл напрямую из [`CompiledModule`]:
//! каждая непустая секция модуля становится своим `PT_LOAD`-сегментом, с
//! правами доступа, зависящими от [`SectionKind`]. `Abs64`-релокации
//! (`.quad my_label`) патчатся здесь, с учётом реальной базы загрузки —
//! раньше это патчилось в кодогенераторе с допущением `base = 0`, что было
//! корректно только для JIT, но не для AOT.
//!
//! ## Структура файла, который мы генерируем
//!
//! ```text
//! Offset  Size   Что
//! ──────  ─────  ──────────────────────────────────────────
//! 0x00    64     ELF Header
//! 0x40    56×N   Program Headers (один на каждую непустую секцию)
//! ...     ...    Байты секций (в порядке их появления в module.sections)
//! ```
//!
//! Это **статически слинкованный** исполняемый файл — никакой динамической линковки.
//! Точка входа берётся из [`CompiledModule::entry`] (метка `_start` или `main`).
//!
//! ## Адреса загрузки
//!
//! По умолчанию код загружается по виртуальному адресу `0x400000` (стандарт для x86_64 Linux).
//! Это соответствует тому, что делает `ld` по умолчанию.
use std::io::{self, Write};

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

/// Базовый виртуальный адрес загрузки кода. Стандарт для x86_64 Linux.
const LOAD_BASE: u64 = 0x40_0000;

/// Адрес, с которого начинается первый `PT_LOAD` в файле (после ELF header +
/// program headers). Выровнен на страницу, чтобы `p_offset` и `p_vaddr`
/// совпадали по модулю страницы, как того требует ELF (для `mod p_align`).
const FIRST_LOAD_FILE_OFFSET: u64 = 0x1000;

/// Возвращает права доступа `PT_LOAD`-сегмента (`PF_R | PF_W | PF_X`) для
/// данного вида секции.
fn segment_flags(kind: SectionKind) -> u32 {
    match kind {
        SectionKind::Text => 5,                    // PF_R | PF_X
        SectionKind::Data | SectionKind::Bss => 6, // PF_R | PF_W
        SectionKind::Custom => 6,
    }
}

/// # ELF Writer
///
/// Принимает [`CompiledModule`] целиком и записывает готовый ELF64 в любой `Write`.
pub struct ElfWriter<'a> {
    module: &'a CompiledModule,
}

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

    /// Записывает готовый ELF64 в переданный writer (файл, Vec<u8>, и т.д.).
    pub fn write<W: Write>(&self, out: &mut W) -> io::Result<()> {
        let layout = Layout::build(self.module);
        let patched = layout.patch_relocations(self.module);

        let phnum = layout.segments.len().max(1);
        let entry_vaddr = layout.entry_vaddr(self.module);

        out.write_all(&self.elf_header(entry_vaddr, phnum as u16))?;

        for seg in &layout.segments {
            out.write_all(&self.phdr(seg))?;
        }
        if layout.segments.is_empty() {
            // ELF без секций — невалидная программа, но не падаем: пишем
            // нулевой phdr, чтобы файл хотя бы имел корректную структуру.
            out.write_all(&[0u8; 56])?;
        }

        let mut written: u64 = 64 + 56 * phnum as u64;
        for seg in &layout.segments {
            pad_to(out, &mut written, seg.file_offset)?;
            if seg.file_size == 0 {
                continue;
            }
            let bytes = &patched[seg.section.0];
            out.write_all(bytes)?;
            written += seg.file_size;
        }

        Ok(())
    }

    // =========================================================================
    // ELF Header (64 байта)
    // =========================================================================

    fn elf_header(&self, entry_vaddr: u64, phnum: u16) -> [u8; 64] {
        let mut h = [0u8; 64];

        // e_ident: magic, class, data, version, OS/ABI
        h[0..4].copy_from_slice(b"\x7fELF");
        h[4] = 2; // ELFCLASS64
        h[5] = 1; // ELFDATA2LSB (little-endian)
        h[6] = 1; // EV_CURRENT
        h[7] = 0; // ELFOSABI_NONE (System V)

        h[16..18].copy_from_slice(&2u16.to_le_bytes()); // e_type = ET_EXEC
        h[18..20].copy_from_slice(&62u16.to_le_bytes()); // e_machine = EM_X86_64
        h[20..24].copy_from_slice(&1u32.to_le_bytes()); // e_version
        h[24..32].copy_from_slice(&entry_vaddr.to_le_bytes()); // e_entry
        h[32..40].copy_from_slice(&64u64.to_le_bytes()); // e_phoff
        h[52..54].copy_from_slice(&64u16.to_le_bytes()); // e_ehsize
        h[54..56].copy_from_slice(&56u16.to_le_bytes()); // e_phentsize
        h[56..58].copy_from_slice(&phnum.to_le_bytes()); // e_phnum
        h[58..60].copy_from_slice(&64u16.to_le_bytes()); // e_shentsize
        // e_shnum = 0, e_shstrndx = 0 — секционных заголовков нет (минимальный ELF)

        h
    }

    // =========================================================================
    // Program Header (56 байт)
    // =========================================================================

    #[allow(clippy::too_many_arguments)]
    fn phdr_raw(
        &self,
        p_type: u32,
        p_flags: u32,
        p_offset: u64,
        p_vaddr: u64,
        p_filesz: u64,
        p_memsz: u64,
        p_align: u64,
    ) -> [u8; 56] {
        let mut p = [0u8; 56];
        p[0..4].copy_from_slice(&p_type.to_le_bytes());
        p[4..8].copy_from_slice(&p_flags.to_le_bytes());
        p[8..16].copy_from_slice(&p_offset.to_le_bytes());
        p[16..24].copy_from_slice(&p_vaddr.to_le_bytes());
        p[24..32].copy_from_slice(&p_vaddr.to_le_bytes()); // p_paddr = p_vaddr
        p[32..40].copy_from_slice(&p_filesz.to_le_bytes());
        p[40..48].copy_from_slice(&p_memsz.to_le_bytes());
        p[48..56].copy_from_slice(&p_align.to_le_bytes());
        p
    }

    fn phdr(&self, seg: &Segment) -> [u8; 56] {
        self.phdr_raw(
            1, // PT_LOAD
            seg.flags,
            seg.file_offset,
            seg.vaddr,
            seg.file_size,
            seg.mem_size,
            0x1000,
        )
    }
}

// =========================================================================
// Layout — вычисляет file offset / vaddr для каждой непустой секции модуля
// =========================================================================

struct Segment {
    section: SectionId,
    flags: u32,
    file_offset: u64,
    vaddr: u64,
    /// Сколько байт реально лежит в файле (0 для `.bss`).
    file_size: u64,
    /// Сколько байт нужно зарезервировать в памяти (>= file_size).
    mem_size: u64,
}

struct Layout {
    segments: Vec<Segment>,
}

impl Layout {
    /// Строит layout: один `PT_LOAD` на каждую секцию модуля с ненулевым
    /// размером (по `virtual_size`, не только по факту записанных байт —
    /// так `.bss` получает место в памяти, не имея байт в файле).
    fn build(module: &CompiledModule) -> Self {
        let mut segments = Vec::new();
        let mut file_offset = FIRST_LOAD_FILE_OFFSET;
        let mut vaddr = LOAD_BASE + FIRST_LOAD_FILE_OFFSET;

        for section in &module.sections {
            if section.virtual_size == 0 {
                continue;
            }

            let is_bss = section.kind == SectionKind::Bss;
            let file_size = if is_bss {
                0
            } else {
                section.bytes.len() as u64
            };
            let mem_size = section.virtual_size as u64;

            segments.push(Segment {
                section: section.id,
                flags: segment_flags(section.kind),
                file_offset,
                vaddr,
                file_size,
                mem_size,
            });

            // Следующий сегмент начинается сразу после этого, выровненный
            // на страницу — так vaddr и file_offset остаются согласованы
            // по модулю 0x1000, как того требует ELF для PT_LOAD.
            file_offset = align_up(file_offset + file_size, 0x1000);
            vaddr = align_up(vaddr + mem_size, 0x1000);
        }

        Self { segments }
    }

    fn segment_for(&self, section: SectionId) -> Option<&Segment> {
        self.segments.iter().find(|s| s.section == section)
    }

    /// Виртуальный адрес точки входа: `vaddr(entry.section) + entry.offset`.
    /// Если точка входа не найдена (`module.entry` пуст), возвращает
    /// `LOAD_BASE` — невалидный, но не паникующий запасной вариант.
    fn entry_vaddr(&self, module: &CompiledModule) -> u64 {
        module
            .entry
            .as_ref()
            .and_then(|entry| {
                self.segment_for(entry.section)
                    .map(|seg| seg.vaddr + entry.offset as u64)
            })
            .unwrap_or(LOAD_BASE)
    }

    /// Возвращает финальный виртуальный адрес символа `name`, если он
    /// определён (имеет секцию) в модуле.
    fn symbol_vaddr(&self, module: &CompiledModule, name: &str) -> Option<u64> {
        let symbol = module.symbol(name)?;
        let section = symbol.section?;
        let seg = self.segment_for(section)?;
        Some(seg.vaddr + symbol.offset as u64)
    }

    /// Возвращает копии байт каждой секции модуля (индексированные по
    /// `SectionId.0`), с патченными `Abs64`-релокациями.
    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 {
                // Неизвестный символ — релокация останется как заглушка
                // (нулевые байты). Резолвинг меток уже проверяется на этапе
                // кодогенерации (`UndefinedLabel`); если мы здесь, значит
                // это какой-то другой случай (например, внешний символ без
                // адреса) — записывать в ELF мы такое пока не умеем.
                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
    }
}

/// Дописывает нули в `out`, пока `*written` не достигнет `target`.
/// Используется для выравнивания файлового offset между сегментами.
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 align_up(val: u64, align: u64) -> u64 {
    (val + align - 1) & !(align - 1)
}
