// rasm-core/src/codegen/aot/pe.rs
//
//! # PE32+ Writer
//!
//! Собирает минимальный исполняемый PE32+ (.exe) для Windows x86_64.
//!
//! ## Структура файла
//!
//! ```text
//! 0x000   64   DOS Header (MZ)
//! 0x040   64   DOS Stub
//! 0x080    4   PE Signature ("PE\0\0")
//! 0x084   20   COFF Header
//! 0x098  240   Optional Header PE32+ (включая 16 Data Directories)
//! 0x188   40   Section Header: .text
//! 0x1B0   40   Section Header: .idata
//! 0x200   N    .text (пользовательский код + wrapper)
//! aligned M    .idata (Import Table + IAT + имена)
//! ```
//!
//! ## Wrapper точки входа
//!
//! В конец `.text` добавляется 19-байтовый wrapper:
//! ```x86asm
//! call  <user_entry>       ; вызываем код пользователя
//! mov   rcx, rax           ; rax -> rcx (первый аргумент ExitProcess)
//! call  [rip + <iat_rel>]  ; ExitProcess(exit_code) через IAT
//! int3                     ; страховка
//! ```
//!
//! ## Импорт
//!
//! Импортируется одна функция: `ExitProcess` из `kernel32.dll`.
use std::io::{self, Write};

// =========================================================================
// Константы
// =========================================================================

const IMAGE_BASE: u64 = 0x0040_0000;
const SECTION_ALIGNMENT: u32 = 0x1000;
const FILE_ALIGNMENT: u32 = 0x200;
const IMAGE_FILE_MACHINE_AMD64: u16 = 0x8664;

const IMAGE_FILE_EXECUTABLE_IMAGE: u16 = 0x0002;
const IMAGE_FILE_LARGE_ADDRESS_AWARE: u16 = 0x0020;

const IMAGE_SCN_CNT_CODE: u32 = 0x0000_0020;
const IMAGE_SCN_CNT_INITIALIZED_DATA: u32 = 0x0000_0040;
const IMAGE_SCN_MEM_EXECUTE: u32 = 0x2000_0000;
const IMAGE_SCN_MEM_READ: u32 = 0x4000_0000;
const IMAGE_SCN_MEM_WRITE: u32 = 0x8000_0000;

const IMAGE_SUBSYSTEM_WINDOWS_CUI: u16 = 3;
const IMAGE_DLLCHARACTERISTICS_NX_COMPAT: u16 = 0x0100;
const IMAGE_DLLCHARACTERISTICS_NO_SEH: u16 = 0x0400;
const IMAGE_DLLCHARACTERISTICS_TERMINAL_SERVER_AWARE: u16 = 0x8000;

// =========================================================================
// DOS Stub
// =========================================================================

const DOS_STUB: [u8; 64] = [
    0x0E, 0x1F, 0xBA, 0x0E, 0x00, 0xB4, 0x09, 0xCD, 0x21, 0xB8, 0x01, 0x4C, 0xCD, 0x21, 0x54, 0x68,
    0x69, 0x73, 0x20, 0x70, 0x72, 0x6F, 0x67, 0x72, 0x61, 0x6D, 0x20, 0x63, 0x61, 0x6E, 0x6E, 0x6F,
    0x74, 0x20, 0x62, 0x65, 0x20, 0x72, 0x75, 0x6E, 0x20, 0x69, 0x6E, 0x20, 0x44, 0x4F, 0x53, 0x20,
    0x6D, 0x6F, 0x64, 0x65, 0x2E, 0x0D, 0x0D, 0x0A, 0x24, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
];

// =========================================================================
// PeWriter
// =========================================================================

pub struct PeWriter {
    pub text: Vec<u8>,
    pub entry_offset: usize,
}

impl PeWriter {
    pub fn new(text: Vec<u8>, entry_offset: usize) -> Self {
        Self { text, entry_offset }
    }

    /// Записывает готовый PE32+ в переданный writer.
    pub fn write<W: Write>(&self, out: &mut W) -> io::Result<()> {
        let ctx = PeContext::new(&self.text, self.entry_offset);
        ctx.emit(out, &self.text, self.entry_offset)
    }
}

// =========================================================================
// PeContext — все вычисленные смещения и размеры
// =========================================================================

/// Все смещения и размеры, вычисленные один раз в `new()`.
struct PeContext {
    text_file_offset: u32,
    idata_file_offset: u32,

    text_virtual_size: u32,
    text_raw_size: u32,
    idata_virtual_size: u32,
    idata_raw_size: u32,

    text_rva: u32,
    idata_rva: u32,
    /// RVA wrapper-а (точка входа .exe).
    entry_rva: u32,
    /// RVA Import Directory (первый IMAGE_IMPORT_DESCRIPTOR).
    import_dir_rva: u32,
    /// RVA IAT (ячейка с адресом ExitProcess).
    iat_rva: u32,

    idata: Vec<u8>,
}

impl PeContext {
    fn new(user_text: &[u8], _user_entry: usize) -> Self {
        // Размер заголовков в файле:
        // DOS(64) + Stub(64) + PESig(4) + COFF(20) + OptHdr(240) + 2×SectHdr(80)
        let headers_size: u32 = 64 + 64 + 4 + 20 + 240 + 80;
        let headers_file_size = align_up(headers_size, FILE_ALIGNMENT);

        // .text: пользовательский код + wrapper (19 байт)
        let wrapper_size: u32 = 19;
        let text_virtual_size = user_text.len() as u32 + wrapper_size;
        let text_raw_size = align_up(text_virtual_size, FILE_ALIGNMENT);
        let text_file_offset = headers_file_size;
        let text_rva = align_up(headers_file_size, SECTION_ALIGNMENT);

        // .idata:
        // [0..20]  IMAGE_IMPORT_DESCRIPTOR kernel32.dll
        // [20..40] IMAGE_IMPORT_DESCRIPTOR sentinel (нули)
        // [40..48] IAT[0] = RVA Hint/Name (u64)
        // [48..56] IAT[1] = 0 (u64)
        // [56..64] INT[0] = RVA Hint/Name (u64)
        // [64..72] INT[1] = 0 (u64)
        // [72..86] Hint(2) + "ExitProcess\0" (14)
        // [86..99] "kernel32.dll\0" (13)
        let idata_virtual_size: u32 = 99;
        let idata_raw_size = align_up(idata_virtual_size, FILE_ALIGNMENT);
        let idata_file_offset = text_file_offset + text_raw_size;
        let idata_rva = align_up(
            text_rva + align_up(text_virtual_size, SECTION_ALIGNMENT),
            SECTION_ALIGNMENT,
        );

        // Смещения внутри .idata
        let import_dir_rva = idata_rva; // offset 0
        let iat_rva = idata_rva + 40; // offset 40
        let int_rva = idata_rva + 56; // offset 56
        let name_rva = idata_rva + 72; // offset 72 (Hint+Name)
        let dll_name_rva = idata_rva + 86; // offset 86

        // Строим байты .idata
        let mut idata = vec![0u8; idata_virtual_size as usize];
        // Import Directory[0]
        put_u32(&mut idata, 0, int_rva); // OriginalFirstThunk
        put_u32(&mut idata, 4, 0); // TimeDateStamp
        put_u32(&mut idata, 8, 0); // ForwarderChain
        put_u32(&mut idata, 12, dll_name_rva); // Name RVA
        put_u32(&mut idata, 16, iat_rva); // FirstThunk (IAT)
        // [20..40] sentinel — уже нули
        // IAT[0]
        put_u64(&mut idata, 40, name_rva as u64);
        // INT[0]
        put_u64(&mut idata, 56, name_rva as u64);
        // Hint/Name
        put_u16(&mut idata, 72, 0);
        idata[74..86].copy_from_slice(b"ExitProcess\0");
        // DLL name
        idata[86..99].copy_from_slice(b"kernel32.dll\0");

        // RVA wrapper-а = text_rva + len(user_text)
        let entry_rva = text_rva + user_text.len() as u32;

        Self {
            text_file_offset,
            idata_file_offset,
            text_virtual_size,
            text_raw_size,
            idata_virtual_size,
            idata_raw_size,
            text_rva,
            idata_rva,
            entry_rva,
            import_dir_rva,
            iat_rva,
            idata,
        }
    }

    /// Строит 19-байтовый wrapper с правильными rel32.
    fn build_wrapper(&self, user_entry_offset: usize) -> [u8; 19] {
        let mut w = [0u8; 19];

        // E8 rel32 — CALL user_entry (5 байт)
        let user_entry_rva = self.text_rva + user_entry_offset as u32;
        let call_next_rva = self.entry_rva + 5;
        let call_rel32 = (user_entry_rva as i64 - call_next_rva as i64) as i32;
        w[0] = 0xE8;
        w[1..5].copy_from_slice(&call_rel32.to_le_bytes());

        // 48 89 C1 — MOV rcx, rax (3 байта)
        w[5] = 0x48;
        w[6] = 0x89;
        w[7] = 0xC1;

        // 48 83 EC 28 — SUB rsp, 40 (4 байта) -> Выравнивание + Shadow Space
        w[8] = 0x48;
        w[9] = 0x83;
        w[10] = 0xEC;
        w[11] = 0x28;

        // FF 15 rel32 — CALL [RIP + rel32] (ExitProcess через IAT) (6 байт)
        // Инструкция заканчивается на: entry_rva + 5 (call) + 3 (mov) + 4 (sub) + 6 (this call) = entry_rva + 18
        let iat_call_next = self.entry_rva + 18;
        let iat_rel32 = (self.iat_rva as i64 - iat_call_next as i64) as i32;
        w[12] = 0xFF;
        w[13] = 0x15; // Смещено с индексов 8..9 на 12..13
        w[14..18].copy_from_slice(&iat_rel32.to_le_bytes());

        // CC — INT3 (1 байт)
        w[18] = 0xCC;

        w
    }

    /// Записывает весь PE-файл.
    fn emit<W: Write>(
        &self,
        out: &mut W,
        user_text: &[u8],
        user_entry_offset: usize,
    ) -> io::Result<()> {
        let wrapper = self.build_wrapper(user_entry_offset);

        out.write_all(&self.dos_header())?;
        out.write_all(&DOS_STUB)?;
        out.write_all(b"PE\0\0")?;
        out.write_all(&self.coff_header())?;
        out.write_all(&self.optional_header())?;
        out.write_all(&self.section_header_text())?;
        out.write_all(&self.section_header_idata())?;

        let headers_written: u32 = 64 + 64 + 4 + 20 + 240 + 80;
        pad(out, headers_written, self.text_file_offset)?;

        out.write_all(user_text)?;
        out.write_all(&wrapper)?;
        pad(out, user_text.len() as u32 + 19, self.text_raw_size)?;

        out.write_all(&self.idata)?;
        pad(out, self.idata_virtual_size, self.idata_raw_size)?;

        Ok(())
    }

    // =========================================================================
    // Заголовки
    // =========================================================================

    fn dos_header(&self) -> [u8; 64] {
        let mut h = [0u8; 64];
        h[0] = b'M';
        h[1] = b'Z';
        // e_lfanew = 0x80 (сразу после DOS Header + Stub)
        put_u32_arr(&mut h, 60, 0x80);
        h
    }

    fn coff_header(&self) -> [u8; 20] {
        let mut h = [0u8; 20];
        put_u16_arr(&mut h, 0, IMAGE_FILE_MACHINE_AMD64);
        put_u16_arr(&mut h, 2, 2); // NumberOfSections
        put_u16_arr(&mut h, 16, 240); // SizeOfOptionalHeader
        put_u16_arr(
            &mut h,
            18,
            IMAGE_FILE_EXECUTABLE_IMAGE | IMAGE_FILE_LARGE_ADDRESS_AWARE,
        );
        h
    }

    /// Optional Header PE32+ (240 байт = 112 стандартных + 128 Data Directories).
    fn optional_header(&self) -> [u8; 240] {
        let mut h = [0u8; 240];

        put_u16_arr(&mut h, 0, 0x020B); // Magic: PE32+
        h[2] = 14; // MajorLinkerVersion

        put_u32_arr(&mut h, 4, self.text_raw_size); // SizeOfCode
        put_u32_arr(&mut h, 8, self.idata_raw_size); // SizeOfInitializedData
        put_u32_arr(&mut h, 16, self.entry_rva); // AddressOfEntryPoint
        put_u32_arr(&mut h, 20, self.text_rva); // BaseOfCode
        put_u64_arr(&mut h, 24, IMAGE_BASE); // ImageBase

        put_u32_arr(&mut h, 32, SECTION_ALIGNMENT);
        put_u32_arr(&mut h, 36, FILE_ALIGNMENT);

        put_u16_arr(&mut h, 40, 6); // MajorOSVersion (Vista+)
        put_u16_arr(&mut h, 48, 6); // MajorSubsystemVersion

        let size_of_image = align_up(
            self.idata_rva + align_up(self.idata_virtual_size, SECTION_ALIGNMENT),
            SECTION_ALIGNMENT,
        );
        put_u32_arr(&mut h, 56, size_of_image);

        let headers_size = align_up(64 + 64 + 4 + 20 + 240 + 80, FILE_ALIGNMENT);
        put_u32_arr(&mut h, 60, headers_size); // SizeOfHeaders

        put_u16_arr(&mut h, 68, IMAGE_SUBSYSTEM_WINDOWS_CUI);
        put_u16_arr(
            &mut h,
            70,
            IMAGE_DLLCHARACTERISTICS_NX_COMPAT
                | IMAGE_DLLCHARACTERISTICS_NO_SEH
                | IMAGE_DLLCHARACTERISTICS_TERMINAL_SERVER_AWARE,
        );

        put_u64_arr(&mut h, 72, 0x10_0000); // SizeOfStackReserve
        put_u64_arr(&mut h, 80, 0x1000); // SizeOfStackCommit
        put_u64_arr(&mut h, 88, 0x10_0000); // SizeOfHeapReserve
        put_u64_arr(&mut h, 96, 0x1000); // SizeOfHeapCommit

        put_u32_arr(&mut h, 108, 16); // NumberOfRvaAndSizes

        // Data Directories начинаются с offset 112:
        // [1] Import Table
        put_u32_arr(&mut h, 112 + 8, self.import_dir_rva); // VirtualAddress
        put_u32_arr(&mut h, 112 + 8 + 4, 40); // Size (2 дескриптора)
        // [12] IAT
        put_u32_arr(&mut h, 112 + 12 * 8, self.iat_rva);
        put_u32_arr(&mut h, 112 + 12 * 8 + 4, 16); // Size (2×u64)

        h
    }

    fn section_header_text(&self) -> [u8; 40] {
        make_section_header(
            b".text\0\0\0",
            self.text_virtual_size,
            self.text_rva,
            self.text_raw_size,
            self.text_file_offset,
            IMAGE_SCN_CNT_CODE | IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ,
        )
    }

    fn section_header_idata(&self) -> [u8; 40] {
        make_section_header(
            b".idata\0\0",
            self.idata_virtual_size,
            self.idata_rva,
            self.idata_raw_size,
            self.idata_file_offset,
            IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE,
        )
    }
}

// =========================================================================
// Вспомогательные функции
// =========================================================================

fn make_section_header(
    name: &[u8; 8],
    virtual_size: u32,
    virtual_address: u32,
    raw_size: u32,
    raw_offset: u32,
    characteristics: u32,
) -> [u8; 40] {
    let mut h = [0u8; 40];
    h[0..8].copy_from_slice(name);
    put_u32_arr(&mut h, 8, virtual_size);
    put_u32_arr(&mut h, 12, virtual_address);
    put_u32_arr(&mut h, 16, raw_size);
    put_u32_arr(&mut h, 20, raw_offset);
    put_u32_arr(&mut h, 36, characteristics);
    h
}

fn pad<W: Write>(out: &mut W, written: u32, total: u32) -> io::Result<()> {
    let rem = written % total.max(1);
    if rem == 0 {
        return Ok(());
    }
    let zeros = vec![0u8; (total - rem) as usize];
    out.write_all(&zeros)
}

fn align_up(v: u32, a: u32) -> u32 {
    (v + a - 1) & !(a - 1)
}

fn put_u16(b: &mut [u8], o: usize, v: u16) {
    b[o..o + 2].copy_from_slice(&v.to_le_bytes());
}
fn put_u32(b: &mut [u8], o: usize, v: u32) {
    b[o..o + 4].copy_from_slice(&v.to_le_bytes());
}
fn put_u64(b: &mut [u8], o: usize, v: u64) {
    b[o..o + 8].copy_from_slice(&v.to_le_bytes());
}

fn put_u16_arr(b: &mut [u8], o: usize, v: u16) {
    put_u16(b, o, v);
}
fn put_u32_arr(b: &mut [u8], o: usize, v: u32) {
    put_u32(b, o, v);
}
fn put_u64_arr(b: &mut [u8], o: usize, v: u64) {
    put_u64(b, o, v);
}
