// rasm-core/src/codegen/aot/pe.rs
//
//! # PE32+ Writer
//!
//! Собирает минимальный исполняемый PE32+ (.exe) для Windows x86_64 напрямую
//! из [`CompiledModule`]: каждая непустая секция модуля (`.text`, `.data`,
//! `.bss`, ...) становится своим PE Section Header, плюс служебная `.idata`
//! для импорта `ExitProcess`. `Abs64`-релокации (`.quad my_label`) патчатся
//! здесь, с учётом реального `ImageBase` — раньше это патчилось в
//! кодогенераторе с допущением `base = 0`, что не соответствовало
//! реальному адресу загрузки PE-образа.
//!
//! ## Структура файла
//!
//! ```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)
//! ...    40×N   Section Headers: одна на каждую непустую секцию модуля + .idata
//! ...     ...   Байты секций модуля (в порядке module.sections), потом .idata
//! ```
//!
//! ## Wrapper точки входа
//!
//! В конец `.text` добавляется 19-байтовый wrapper:
//! ```x86asm
//! call  <user_entry>       ; вызываем код пользователя
//! mov   rcx, rax           ; rax -> rcx (первый аргумент ExitProcess)
//! sub   rsp, 40            ; выравнивание стека + shadow space (Win64 ABI)
//! call  [rip + <iat_rel>]  ; ExitProcess(exit_code) через IAT
//! int3                     ; страховка
//! ```
//!
//! ## Импорт
//!
//! Импортируется одна функция: `ExitProcess` из `kernel32.dll`.
use std::{
    collections::{BTreeMap, HashMap},
    io::{self, Write},
};

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

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

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_CNT_UNINITIALIZED_DATA: u32 = 0x0000_0080;
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;

const WRAPPER_SIZE: u32 = 19;
const IMPORT_THUNK_SIZE: u32 = 6;

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

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

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

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

// =========================================================================
// Раскладка одной пользовательской секции внутри PE
// =========================================================================

struct PeSection {
    /// Индекс в `module.sections` (для поиска байт/релокаций).
    module_index: usize,
    name: [u8; 8],
    characteristics: u32,
    file_offset: u32,
    /// Сколько байт реально лежит в файле (0 для `.bss`).
    raw_size: u32,
    rva: u32,
    /// Сколько байт нужно зарезервировать в памяти (>= raw_size для `.bss`).
    virtual_size: u32,
}

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

fn pe_characteristics(kind: SectionKind) -> u32 {
    match kind {
        SectionKind::Text => IMAGE_SCN_CNT_CODE | IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ,
        SectionKind::Data | SectionKind::Custom => {
            IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE
        }
        SectionKind::Bss => {
            IMAGE_SCN_CNT_UNINITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE
        }
    }
}

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

/// Все смещения и размеры, вычисленные один раз в `new()`.
#[derive(Debug, Clone)]
struct PeImport {
    module: String,
    name: String,
    alias: String,
}

fn collect_imports(module: &CompiledModule) -> Vec<PeImport> {
    let mut imports: Vec<_> = module
        .imports
        .iter()
        .map(|import| PeImport {
            module: import.module.clone(),
            name: import.name.clone(),
            alias: import.alias.clone(),
        })
        .collect();

    imports.push(PeImport {
        module: "kernel32.dll".to_string(),
        name: "ExitProcess".to_string(),
        alias: "ExitProcess".to_string(),
    });

    imports.sort_by(|a, b| {
        a.module
            .to_ascii_lowercase()
            .cmp(&b.module.to_ascii_lowercase())
            .then_with(|| a.name.cmp(&b.name))
            .then_with(|| a.alias.cmp(&b.alias))
    });
    imports.dedup_by(|a, b| {
        a.module.eq_ignore_ascii_case(&b.module) && a.name == b.name && a.alias == b.alias
    });
    imports
}

struct ImportImage {
    bytes: Vec<u8>,
    import_dir_rva: u32,
    import_dir_size: u32,
    iat_rva: u32,
    iat_size: u32,
    iat_by_alias: HashMap<String, u32>,
}

fn build_import_image(imports: &[PeImport], idata_rva: u32) -> ImportImage {
    let mut by_module: BTreeMap<String, Vec<&PeImport>> = BTreeMap::new();
    for import in imports {
        by_module
            .entry(import.module.to_ascii_lowercase())
            .or_default()
            .push(import);
    }

    for group in by_module.values_mut() {
        group.sort_by(|a, b| a.name.cmp(&b.name).then_with(|| a.alias.cmp(&b.alias)));
    }

    let descriptor_count = by_module.len() + 1;
    let descriptors_size = descriptor_count * 20;
    let mut cursor = descriptors_size;

    let mut dll_layouts = Vec::new();
    for (module_key, group) in &by_module {
        let iat_offset = cursor;
        cursor += (group.len() + 1) * 8;
        let int_offset = cursor;
        cursor += (group.len() + 1) * 8;
        dll_layouts.push((module_key.clone(), iat_offset, int_offset));
    }

    let mut name_offsets: HashMap<(String, String), usize> = HashMap::new();
    for (module_key, group) in &by_module {
        for import in group {
            name_offsets.insert((module_key.clone(), import.name.clone()), cursor);
            cursor += 2 + import.name.len() + 1;
            if !cursor.is_multiple_of(2) {
                cursor += 1;
            }
        }
    }

    let mut dll_name_offsets = HashMap::new();
    for (module_key, group) in &by_module {
        dll_name_offsets.insert(module_key.clone(), cursor);
        cursor += group[0].module.len() + 1;
    }

    let mut bytes = vec![0u8; cursor];
    let mut iat_by_alias = HashMap::new();
    let iat_rva = idata_rva + descriptors_size as u32;
    let mut iat_size = 0u32;

    for (descriptor_idx, (module_key, iat_offset, int_offset)) in dll_layouts.iter().enumerate() {
        let group = &by_module[module_key];
        let descriptor_offset = descriptor_idx * 20;
        let dll_name_rva = idata_rva + dll_name_offsets[module_key] as u32;

        put_u32(
            &mut bytes,
            descriptor_offset,
            idata_rva + *int_offset as u32,
        );
        put_u32(&mut bytes, descriptor_offset + 4, 0);
        put_u32(&mut bytes, descriptor_offset + 8, 0);
        put_u32(&mut bytes, descriptor_offset + 12, dll_name_rva);
        put_u32(
            &mut bytes,
            descriptor_offset + 16,
            idata_rva + *iat_offset as u32,
        );

        iat_size += ((group.len() + 1) * 8) as u32;

        for (idx, import) in group.iter().enumerate() {
            let name_key = (module_key.clone(), import.name.clone());
            let name_offset = name_offsets[&name_key];
            let hint_name_rva = idata_rva + name_offset as u32;

            put_u64(&mut bytes, iat_offset + idx * 8, hint_name_rva as u64);
            put_u64(&mut bytes, int_offset + idx * 8, hint_name_rva as u64);
            iat_by_alias.insert(
                import.alias.clone(),
                idata_rva + *iat_offset as u32 + idx as u32 * 8,
            );

            put_u16(&mut bytes, name_offset, 0);
            bytes[name_offset + 2..name_offset + 2 + import.name.len()]
                .copy_from_slice(import.name.as_bytes());
        }

        let dll_name_offset = dll_name_offsets[module_key];
        let module_name = &group[0].module;
        bytes[dll_name_offset..dll_name_offset + module_name.len()]
            .copy_from_slice(module_name.as_bytes());
    }

    ImportImage {
        bytes,
        import_dir_rva: idata_rva,
        import_dir_size: descriptors_size as u32,
        iat_rva,
        iat_size,
        iat_by_alias,
    }
}

struct PeContext {
    sections: Vec<PeSection>,

    idata_file_offset: u32,
    idata_virtual_size: u32,
    idata_raw_size: 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,
    import_dir_size: u32,
    iat_size: u32,
    iat_by_alias: HashMap<String, u32>,
    thunk_rvas: HashMap<String, u32>,
    thunk_bytes: Vec<u8>,

    idata: Vec<u8>,

    headers_file_size: u32,
}

impl PeContext {
    fn new(module: &CompiledModule) -> Self {
        let imports = collect_imports(module);
        let thunk_bytes_len = imports.len() as u32 * IMPORT_THUNK_SIZE;
        let num_user_sections = module
            .sections
            .iter()
            .filter(|s| s.virtual_size > 0)
            .count();
        let total_sections = num_user_sections + 1; // + .idata

        // Размер заголовков в файле:
        // DOS(64) + Stub(64) + PESig(4) + COFF(20) + OptHdr(240) + N×SectHdr(40)
        let headers_size: u32 = 64 + 64 + 4 + 20 + 240 + 40 * total_sections as u32;
        let headers_file_size = align_up(headers_size, FILE_ALIGNMENT);

        // ---- Раскладываем пользовательские секции модуля ----
        let mut sections = Vec::with_capacity(num_user_sections);
        let mut file_offset = headers_file_size;
        let mut rva = align_up(headers_file_size, SECTION_ALIGNMENT);
        let mut text_entry_rva = None;

        for (idx, section) in module.sections.iter().enumerate() {
            if section.virtual_size == 0 {
                continue;
            }

            let is_text = section.kind == SectionKind::Text;
            let is_bss = section.kind == SectionKind::Bss;

            // Wrapper дописывается сразу после байт пользователя в `.text`.
            let extra = if is_text {
                thunk_bytes_len + WRAPPER_SIZE
            } else {
                0
            };
            let virtual_size = section.virtual_size as u32 + extra;
            let raw_size = if is_bss {
                0
            } else {
                align_up(virtual_size, FILE_ALIGNMENT)
            };

            if is_text {
                // Wrapper идёт сразу за пользовательским кодом этой секции.
                text_entry_rva = Some(rva + section.bytes.len() as u32 + thunk_bytes_len);
            }

            sections.push(PeSection {
                module_index: idx,
                name: pe_section_name(&section.name),
                characteristics: pe_characteristics(section.kind),
                file_offset,
                raw_size,
                rva,
                virtual_size,
            });

            file_offset = align_up(file_offset + raw_size, FILE_ALIGNMENT);
            rva = align_up(
                rva + align_up(virtual_size, SECTION_ALIGNMENT),
                SECTION_ALIGNMENT,
            );
        }

        // ---- .idata (служебная секция импорта ExitProcess) ----
        // [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_file_offset = file_offset;
        let idata_rva = rva;
        let import_image = build_import_image(&imports, idata_rva);
        let idata_virtual_size = import_image.bytes.len() as u32;
        let idata_raw_size = align_up(idata_virtual_size, FILE_ALIGNMENT);

        // Если `.text` не найден среди секций модуля (нет кода вообще) —
        // используем idata_rva как безопасный (хоть и невалидный) fallback,
        // чтобы не паниковать при сборке заголовков.
        let entry_rva = text_entry_rva.unwrap_or(idata_rva);
        let text_rva = sections
            .iter()
            .find(|s| s.characteristics & IMAGE_SCN_MEM_EXECUTE != 0)
            .map(|s| s.rva);
        let text_user_len = module
            .sections
            .iter()
            .find(|s| s.kind == SectionKind::Text)
            .map(|s| s.bytes.len() as u32)
            .unwrap_or(0);
        let mut thunk_rvas = HashMap::new();
        let mut thunk_bytes = Vec::with_capacity(thunk_bytes_len as usize);

        if let Some(text_rva) = text_rva {
            for (idx, import) in imports.iter().enumerate() {
                let thunk_rva = text_rva + text_user_len + idx as u32 * IMPORT_THUNK_SIZE;
                thunk_rvas.insert(import.alias.clone(), thunk_rva);

                let iat_rva = import_image.iat_by_alias[&import.alias];
                let next_rva = thunk_rva + IMPORT_THUNK_SIZE;
                let rel32 = (iat_rva as i64 - next_rva as i64) as i32;
                thunk_bytes.extend_from_slice(&[0xFF, 0x25]);
                thunk_bytes.extend_from_slice(&rel32.to_le_bytes());
            }
        }

        Self {
            sections,
            idata_file_offset,
            idata_virtual_size,
            idata_raw_size,
            idata_rva,
            entry_rva,
            import_dir_rva: import_image.import_dir_rva,
            import_dir_size: import_image.import_dir_size,
            iat_rva: import_image.iat_rva,
            iat_size: import_image.iat_size,
            iat_by_alias: import_image.iat_by_alias,
            thunk_rvas,
            thunk_bytes,
            idata: import_image.bytes,
            headers_file_size,
        }
    }

    fn text_section(&self) -> Option<&PeSection> {
        self.sections.iter().find(|s| {
            // `.text` — единственная секция с флагом EXECUTE среди наших.
            s.characteristics & IMAGE_SCN_MEM_EXECUTE != 0
        })
    }

    /// Строит 19-байтовый wrapper с правильными rel32, используя реальную
    /// точку входа из `module.entry` (если её нет — указывает wrapper на
    /// начало `.text`, что почти наверняка неверно, но не паникует).
    fn build_wrapper(&self, module: &CompiledModule) -> [u8; 19] {
        let mut w = [0u8; 19];

        let text = self.text_section();
        let user_entry_rva = match (&module.entry, text) {
            (Some(entry), Some(text_section)) => text_section.rva + entry.offset as u32,
            (_, Some(text_section)) => text_section.rva,
            (_, None) => 0,
        };

        // E8 rel32 — CALL user_entry (5 байт)
        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 байт)
        let iat_call_next = self.entry_rva + 18;
        let exit_iat_rva = self
            .iat_by_alias
            .get("ExitProcess")
            .copied()
            .unwrap_or(self.iat_rva);
        let iat_rel32 = (exit_iat_rva as i64 - iat_call_next as i64) as i32;
        w[12] = 0xFF;
        w[13] = 0x15;
        w[14..18].copy_from_slice(&iat_rel32.to_le_bytes());

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

        w
    }

    /// Возвращает RVA символа `name`, если он определён в модуле (через
    /// секцию, в которой он лежит, плюс offset внутри неё).
    fn symbol_rva(&self, module: &CompiledModule, name: &str) -> Option<u32> {
        let symbol = module.symbol(name)?;
        let section_id = symbol.section?;
        let pe_section = self
            .sections
            .iter()
            .find(|s| s.module_index == section_id.0)?;
        Some(pe_section.rva + symbol.offset as u32)
    }

    /// Возвращает байты каждой пользовательской секции модуля (в том же
    /// порядке, что `self.sections`), с патченными `Abs64`-релокациями и
    /// дописанным wrapper-ом в конец `.text`.
    fn patch_and_collect(&self, module: &CompiledModule) -> Vec<Vec<u8>> {
        let wrapper = self.build_wrapper(module);

        let mut bytes: Vec<Vec<u8>> = self
            .sections
            .iter()
            .map(|s| module.sections[s.module_index].bytes.clone())
            .collect();

        for reloc in &module.relocations {
            let Some(pe_idx) = self
                .sections
                .iter()
                .position(|s| s.module_index == reloc.section.0)
            else {
                continue;
            };

            let section_rva = self.sections[pe_idx].rva;
            let section_bytes = &mut bytes[pe_idx];
            let offset = reloc.offset;

            match reloc.kind {
                RelocKind::Abs64 => {
                    let target_rva = self
                        .symbol_rva(module, &reloc.target)
                        .or_else(|| self.thunk_rvas.get(&reloc.target).copied());
                    let Some(target_rva) = target_rva else {
                        continue;
                    };
                    let target_vaddr = IMAGE_BASE + target_rva as u64;

                    if offset + 8 <= section_bytes.len() {
                        section_bytes[offset..offset + 8]
                            .copy_from_slice(&target_vaddr.to_le_bytes());
                    }
                }
                RelocKind::Rel32 => {
                    let target_rva = self
                        .thunk_rvas
                        .get(&reloc.target)
                        .copied()
                        .or_else(|| self.symbol_rva(module, &reloc.target));
                    let Some(target_rva) = target_rva else {
                        continue;
                    };
                    let patch_next_rva = section_rva + offset as u32 + 4;
                    let rel32 = (target_rva as i64 - patch_next_rva as i64) as i32;

                    if offset + 4 <= section_bytes.len() {
                        section_bytes[offset..offset + 4].copy_from_slice(&rel32.to_le_bytes());
                    }
                }
            }
        }

        // Дописываем wrapper в конец `.text`.
        if let Some(text_pe_idx) = self
            .sections
            .iter()
            .position(|s| s.characteristics & IMAGE_SCN_MEM_EXECUTE != 0)
        {
            bytes[text_pe_idx].extend_from_slice(&self.thunk_bytes);
            bytes[text_pe_idx].extend_from_slice(&wrapper);
        }

        bytes
    }

    /// Записывает весь PE-файл.
    fn emit<W: Write>(&self, out: &mut W, module: &CompiledModule) -> io::Result<()> {
        let patched = self.patch_and_collect(module);

        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())?;

        for section in &self.sections {
            out.write_all(&self.section_header(section))?;
        }
        out.write_all(&self.idata_section_header())?;

        let mut written: u32 = 64 + 64 + 4 + 20 + 240 + 40 * (self.sections.len() as u32 + 1);

        for (section, bytes) in self.sections.iter().zip(patched.iter()) {
            pad_to(out, &mut written, section.file_offset)?;
            if section.raw_size == 0 {
                continue;
            }
            out.write_all(bytes)?;
            written += bytes.len() as u32;
            pad_to(out, &mut written, section.file_offset + section.raw_size)?;
        }

        pad_to(out, &mut written, self.idata_file_offset)?;
        out.write_all(&self.idata)?;
        written += self.idata_virtual_size;
        pad_to(
            out,
            &mut written,
            self.idata_file_offset + self.idata_raw_size,
        )?;

        Ok(())
    }

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

    fn dos_header(&self) -> [u8; 64] {
        let mut h = [0u8; 64];
        h[0] = b'M';
        h[1] = b'Z';
        put_u32(&mut h, 60, 0x80); // e_lfanew
        h
    }

    fn coff_header(&self) -> [u8; 20] {
        let mut h = [0u8; 20];
        put_u16(&mut h, 0, IMAGE_FILE_MACHINE_AMD64);
        put_u16(&mut h, 2, self.sections.len() as u16 + 1); // NumberOfSections (+ .idata)
        put_u16(&mut h, 16, 240); // SizeOfOptionalHeader
        put_u16(
            &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(&mut h, 0, 0x020B); // Magic: PE32+
        h[2] = 14; // MajorLinkerVersion

        let size_of_code: u32 = self
            .sections
            .iter()
            .filter(|s| s.characteristics & IMAGE_SCN_CNT_CODE != 0)
            .map(|s| s.raw_size)
            .sum();
        let size_of_init_data: u32 = self
            .sections
            .iter()
            .filter(|s| s.characteristics & IMAGE_SCN_CNT_INITIALIZED_DATA != 0)
            .map(|s| s.raw_size)
            .sum::<u32>()
            + self.idata_raw_size;

        put_u32(&mut h, 4, size_of_code); // SizeOfCode
        put_u32(&mut h, 8, size_of_init_data); // SizeOfInitializedData
        put_u32(&mut h, 16, self.entry_rva); // AddressOfEntryPoint
        put_u32(&mut h, 20, self.text_section().map(|s| s.rva).unwrap_or(0)); // BaseOfCode
        put_u64(&mut h, 24, IMAGE_BASE); // ImageBase

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

        put_u16(&mut h, 40, 6); // MajorOSVersion (Vista+)
        put_u16(&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(&mut h, 56, size_of_image);
        put_u32(&mut h, 60, self.headers_file_size); // SizeOfHeaders

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

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

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

        // Data Directories начинаются с offset 112:
        // [1] Import Table
        put_u32(&mut h, 112 + 8, self.import_dir_rva);
        put_u32(&mut h, 112 + 8 + 4, self.import_dir_size);
        // [12] IAT
        put_u32(&mut h, 112 + 12 * 8, self.iat_rva);
        put_u32(&mut h, 112 + 12 * 8 + 4, self.iat_size);

        h
    }

    fn section_header(&self, section: &PeSection) -> [u8; 40] {
        make_section_header(
            &section.name,
            section.virtual_size,
            section.rva,
            section.raw_size,
            section.file_offset,
            section.characteristics,
        )
    }

    fn idata_section_header(&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(&mut h, 8, virtual_size);
    put_u32(&mut h, 12, virtual_address);
    put_u32(&mut h, 16, raw_size);
    put_u32(&mut h, 20, raw_offset);
    put_u32(&mut h, 36, characteristics);
    h
}

/// Дописывает нули в `out`, пока `*written` не достигнет `target`.
fn pad_to<W: Write>(out: &mut W, written: &mut u32, target: u32) -> 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(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());
}
