pe-assembler 0.1.1

use crate::{
    formats::exe::writer::ExeWriter,
    helpers::pe_writer::PeWriter,
    types::{
        tables::{ImportEntry, ImportTable},
        CoffHeader, DataDirectory, DosHeader, NtHeader, OptionalHeader, PeHeader, PeProgram, PeSection, SubsystemType,
    },
};
use gaia_types::{helpers::Architecture, GaiaError};
use std::io::Cursor;
/// PE assembler builder
#[derive(Debug)]
pub struct PeBuilder {
    architecture: Option<Architecture>,
    subsystem: Option<SubsystemType>,
    entry_point: Option<u32>,
    image_base: Option<u64>,
    imports: Vec<(String, Vec<String>)>, // (dll_name, functions)
    code: Option<Vec<u8>>,
    data: Option<Vec<u8>>,
    sections: Vec<PeSection>, // Add this field
}

impl PeBuilder {
    /// Create a new PE assembler builder
    pub fn new() -> Self {
        Self {
            architecture: None,
            subsystem: None,
            entry_point: None,
            image_base: None,
            imports: Vec::new(),
            code: None,
            data: None,
            sections: Vec::new(), // Initialize the new field
        }
    }

    /// Set target architecture
    pub fn architecture(mut self, arch: Architecture) -> Self {
        self.architecture = Some(arch);
        self
    }

    /// Set subsystem type
    pub fn subsystem(mut self, subsystem: SubsystemType) -> Self {
        self.subsystem = Some(subsystem);
        self
    }

    /// Set entry point address
    pub fn entry_point(mut self, entry_point: u32) -> Self {
        self.entry_point = Some(entry_point);
        self
    }

    /// Set image base address
    pub fn image_base(mut self, image_base: u64) -> Self {
        self.image_base = Some(image_base);
        self
    }

    /// Import a single function
    pub fn import_function(mut self, dll_name: &str, function_name: &str) -> Self {
        // Find if the DLL already exists
        if let Some(entry) = self.imports.iter_mut().find(|(name, _)| name == dll_name) {
            entry.1.push(function_name.to_string());
        }
        else {
            self.imports.push((dll_name.to_string(), vec![function_name.to_string()]));
        }
        self
    }

    /// Import multiple functions
    pub fn import_functions(mut self, dll_name: &str, function_names: &[&str]) -> Self {
        let functions: Vec<String> = function_names.iter().map(|&s| s.to_string()).collect();

        // Find if the DLL already exists
        if let Some(entry) = self.imports.iter_mut().find(|(name, _)| name == dll_name) {
            entry.1.extend(functions);
        }
        else {
            self.imports.push((dll_name.to_string(), functions));
        }
        self
    }

    /// Set code data
    pub fn code(mut self, code: Vec<u8>) -> Self {
        self.code = Some(code);
        self
    }

    /// Set data
    pub fn data(mut self, data: Vec<u8>) -> Self {
        self.data = Some(data);
        self
    }

    /// Get import information (for creating import table)
    pub fn get_imports(&self) -> &Vec<(String, Vec<String>)> {
        &self.imports
    }

    /// Generate PE header info
    pub fn build_header(&self) -> Result<PeHeader, GaiaError> {
        let architecture = self
            .architecture
            .as_ref()
            .ok_or_else(|| GaiaError::syntax_error("Architecture is required", gaia_types::SourceLocation::default()))?;
        let pointer_size: u32 = if *architecture == Architecture::X86_64 { 8 } else { 4 };
        let subsystem = self
            .subsystem
            .ok_or_else(|| GaiaError::syntax_error("Subsystem is required", gaia_types::SourceLocation::default()))?;
        let entry_point = self.entry_point.unwrap_or(0x1000);
        let image_base = self.image_base.unwrap_or(match architecture {
            Architecture::X86 => 0x400000,
            Architecture::X86_64 => 0x140000000,
            _ => 0x400000,
        });

        // Create DOS header
        let dos_header = DosHeader::new(0x80);

        // Create NT header
        let nt_header = NtHeader {
            signature: 0x00004550, // "PE\0\0"
        };

        // Create COFF header
        let machine = match architecture {
            Architecture::X86 => 0x014C,
            Architecture::X86_64 => 0x8664,
            _ => 0x014C,
        };

        let mut section_count = 0;
        if self.code.is_some() {
            section_count += 1;
        }
        if self.data.is_some() {
            section_count += 1;
        }
        if !self.imports.is_empty() {
            section_count += 1;
        }

        let optional_header_size = match architecture {
            Architecture::X86_64 => 240,
            _ => 224,
        };

        // Set COFF characteristics based on architecture:
        // - x86: Executable image | 32-bit machine
        // - x64: Executable image | Large address aware (don't set 32-bit machine bit)
        let characteristics = match architecture {
            Architecture::X86 => 0x0102,    // IMAGE_FILE_EXECUTABLE_IMAGE | IMAGE_FILE_32BIT_MACHINE
            Architecture::X86_64 => 0x0022, // IMAGE_FILE_EXECUTABLE_IMAGE | IMAGE_FILE_LARGE_ADDRESS_AWARE
            _ => 0x0102,
        };

        let coff_header = CoffHeader::new(machine, section_count)
            .with_timestamp(0)
            .with_symbol_table(0, 0)
            .with_optional_header_size(optional_header_size)
            .with_characteristics(characteristics);

        // Normalize size_of_code and size_of_initialized_data, using aligned size of sections
        let size_of_code = if let Some(code) = &self.code { ((code.len() + 0x1FF) / 0x200 * 0x200) as u32 } else { 0 };
        let mut size_of_initialized_data = 0;
        if let Some(data) = &self.data {
            size_of_initialized_data += ((data.len() + 0x1FF) / 0x200 * 0x200) as u32;
            // .data
        }
        if !self.imports.is_empty() {
            // We can't accurately predict .idata size here, but 0x200 should be enough for most cases
            // Or we could build sections first to get the exact size
            size_of_initialized_data += 0x200; // .idata
        }

        // Calculate size_of_image: starting from 0x1000, each existing section adds 0x1000
        let mut size_of_image = 0x1000; // DOS/Headers take one alignment page
        if self.code.is_some() {
            size_of_image += 0x1000;
        }
        if self.data.is_some() {
            size_of_image += 0x1000;
        }
        if !self.imports.is_empty() {
            size_of_image += 0x1000;
        }

        let mut optional_header = OptionalHeader::new_for_architecture(
            architecture,
            entry_point,
            image_base,
            size_of_code,
            0x200,         // size_of_headers
            size_of_image, // Dynamically calculated image size
            subsystem,
        );
        optional_header.size_of_initialized_data = size_of_initialized_data;

        // Disable ASLR (DYNAMIC_BASE), otherwise the absolute addresses we patch will be invalid due to base randomization
        // DYNAMIC_BASE bit value is 0x0040
        optional_header.dll_characteristics &= !0x0040;

        // Set import table data directory (dynamically calculate RVA)
        // Compatibility mode (x64): both IAT and INT initially point to the RVA of IMAGE_IMPORT_BY_NAME (Hint+Name).
        // - x64: INT = array of name pointers, IAT = name RVA (overwritten with real address after loader resolution)
        // - x86: OFT/INT = 0, IAT = name RVA (common compatible layout)
        if !self.imports.is_empty() {
            // Find the actual RVA of the .idata section
            let idata_section = self
                .sections
                .iter()
                .find(|s| s.name == ".idata")
                .ok_or_else(|| GaiaError::syntax_error("Missing .idata section", gaia_types::SourceLocation::default()))?;
            let import_rva_base = idata_section.virtual_address;

            // Subsequent calculations remain unchanged, but use import_rva_base instead of hardcoded values
            let mut current_rva = import_rva_base + ((self.imports.len() + 1) as u32) * 20;
            for (dll_name, _) in &self.imports {
                current_rva += (dll_name.len() as u32) + 1;
            }
            if current_rva % 2 != 0 {
                current_rva += 1;
            }
            // Function Hint/Name
            for (_, functions) in &self.imports {
                for func in functions {
                    // Name aligned to 2 bytes
                    if current_rva % 2 != 0 {
                        current_rva += 1;
                    }
                    current_rva += 2 + (func.len() as u32) + 1;
                }
            }
            // Remove redundant alignment here, keep consistent with write_import_table
            // INT
            if current_rva % pointer_size != 0 {
                current_rva = (current_rva + pointer_size - 1) & !(pointer_size - 1);
            }
            for (_, functions) in &self.imports {
                current_rva += ((functions.len() as u32) + 1) * pointer_size;
            }
            // IAT (start of IAT at this point)
            if current_rva % pointer_size != 0 {
                current_rva = (current_rva + pointer_size - 1) & !(pointer_size - 1);
            }
            let iat_rva_start = current_rva; // Record IAT start
                                             // IAT length
            let mut end_rva = current_rva;
            for (_, functions) in &self.imports {
                end_rva += ((functions.len() as u32) + 1) * pointer_size;
            }
            optional_header.data_directories[1] =
                DataDirectory { virtual_address: import_rva_base, size: end_rva - import_rva_base };
            // Also fill IAT Directory (index 12) so loader knows IAT range
            let mut iat_rva_end = iat_rva_start;
            for (_, functions) in &self.imports {
                iat_rva_end += ((functions.len() as u32) + 1) * pointer_size;
            }
            optional_header.data_directories[12] =
                DataDirectory { virtual_address: iat_rva_start, size: iat_rva_end - iat_rva_start };
        }

        Ok(PeHeader { dos_header, nt_header, coff_header, optional_header })
    }

    /// Generate section list
    pub fn build_sections(&mut self) -> Vec<PeSection> {
        let mut sections = Vec::new();
        let mut next_virtual_address = 0x1000;
        let mut next_raw_data_offset = 0x200;

        // Pre-calculate section RVAs for fix_code_relocations
        // Note: Logic here must be exactly consistent with section addition order below
        let mut code_rva = None;
        let mut data_rva = None;

        if self.code.is_some() {
            code_rva = Some(next_virtual_address);
            next_virtual_address += 0x1000;
        }
        if self.data.is_some() {
            data_rva = Some(next_virtual_address);
            next_virtual_address += 0x1000;
        }

        // Reset for actual building
        next_virtual_address = 0x1000;

        // Add code section
        if let Some(code) = &self.code {
            let mut code_data = code.clone();

            // Fix relocations
            self.fix_code_relocations_with_rvas(&mut code_data, code_rva.unwrap_or(0x1000), data_rva);

            // Align to 512 bytes
            let raw_size = ((code_data.len() + 0x1FF) / 0x200 * 0x200) as u32;
            while code_data.len() < raw_size as usize {
                code_data.push(0);
            }

            let text_section = PeSection {
                name: ".text".to_string(),
                virtual_size: 0x1000,
                virtual_address: next_virtual_address,
                size_of_raw_data: raw_size,
                pointer_to_raw_data: next_raw_data_offset,
                pointer_to_relocations: 0,
                pointer_to_line_numbers: 0,
                number_of_relocations: 0,
                number_of_line_numbers: 0,
                characteristics: 0x60000020,
                data: code_data,
            };
            sections.push(text_section);
            next_virtual_address += 0x1000;
            next_raw_data_offset += raw_size;
        }

        // Add data section
        if let Some(data) = &self.data {
            let mut data_bytes = data.clone();
            // Align to 512 bytes
            let raw_size = ((data_bytes.len() + 0x1FF) / 0x200 * 0x200) as u32;
            while data_bytes.len() < raw_size as usize {
                data_bytes.push(0);
            }

            let data_section = PeSection {
                name: ".data".to_string(),
                virtual_size: 0x1000,
                virtual_address: next_virtual_address,
                size_of_raw_data: raw_size,
                pointer_to_raw_data: next_raw_data_offset,
                pointer_to_relocations: 0,
                pointer_to_line_numbers: 0,
                number_of_relocations: 0,
                number_of_line_numbers: 0,
                characteristics: 0xC0000040,
                data: data_bytes,
            };
            sections.push(data_section);
            next_virtual_address += 0x1000;
            next_raw_data_offset += raw_size;
        }

        // Add import table section (if there are imports)
        if !self.imports.is_empty() {
            let mut idata_section = self.build_import_section();
            idata_section.virtual_address = next_virtual_address;
            idata_section.pointer_to_raw_data = next_raw_data_offset;
            sections.push(idata_section);
        }

        sections
    }

    /// Build import table section
    fn build_import_section(&self) -> PeSection {
        // Data is not filled here, let write_import_table handle it.
        // Note: write_import_table uses "compatible mode" (see above), where IAT initially fills in Hint/Name RVA on x64.
        PeSection {
            name: ".idata".to_string(),
            virtual_size: 0x1000,
            virtual_address: 0x3000, // This value will be overwritten in build_sections
            size_of_raw_data: 0x200,
            pointer_to_raw_data: 0x600, // This value will be overwritten in build_sections
            pointer_to_relocations: 0,
            pointer_to_line_numbers: 0,
            number_of_relocations: 0,
            number_of_line_numbers: 0,
            characteristics: 0xC0000040, // IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE
            data: Vec::new(),            // Empty data, filled by write_import_table method
        }
    }

    /// Fix relocations in code
    fn fix_code_relocations_with_rvas(&self, code: &mut Vec<u8>, code_section_rva: u32, data_section_rva: Option<u32>) {
        // Find placeholders for CALL instructions and replace them with correct addresses
        let mut i = 0;

        let arch = self.architecture.as_ref().unwrap_or(&Architecture::X86).clone();
        let pointer_size: usize = if arch == Architecture::X86_64 { 8 } else { 4 };
        let image_base: u64 = self.image_base.as_ref().copied().unwrap_or(match arch {
            Architecture::X86 => 0x400000,
            Architecture::X86_64 => 0x140000000,
            _ => 0x400000,
        });

        // Calculate start RVA of .idata section (consistent with build_sections)
        let mut import_rva_base: u64 = 0x1000; // Starting RVA
        if self.code.is_some() {
            import_rva_base += 0x1000;
        }
        if self.data.is_some() {
            import_rva_base += 0x1000;
        }

        // Calculate actual position of IAT (must be exactly consistent with calculation in write_import_table)
        let mut current_rva: u64 = import_rva_base;

        // Calculate IAT offset (keep consistent with write_import_table)
        if !self.imports.is_empty() {
            // Import descriptor table size: (number of DLLs + 1) * 20 bytes
            current_rva += ((self.imports.len() + 1) * 20) as u64;

            // DLL name size
            for (dll_name, _) in &self.imports {
                current_rva += (dll_name.len() + 1) as u64; // include null terminator
            }

            // Align to 2 bytes, consistent with write_import_table
            if current_rva % 2 != 0 {
                current_rva += 1;
            }

            // Function name (Hint(2) + Name + '\0'), accumulate one by one
            for (_, functions) in &self.imports {
                for function in functions {
                    // Name aligned to 2 bytes
                    if current_rva % 2 != 0 {
                        current_rva += 1;
                    }
                    current_rva += 2 + (function.len() + 1) as u64;
                }
            }

            // INT (OriginalFirstThunk) aligned to pointer_size bytes and allocated
            if current_rva % pointer_size as u64 != 0 {
                current_rva = (current_rva + pointer_size as u64 - 1) & !(pointer_size as u64 - 1);
            }
            for (_, functions) in &self.imports {
                current_rva += ((functions.len() as u64) + 1) * pointer_size as u64;
                // include terminator
            }

            // IAT (FirstThunk) aligned to pointer_size bytes and allocated -> this is the IAT start RVA
            if current_rva % pointer_size as u64 != 0 {
                current_rva = (current_rva + pointer_size as u64 - 1) & !(pointer_size as u64 - 1);
            }
        }

        let iat_start_rva = current_rva;

        while i < code.len() {
            // x64: patch multiple instructions using [rip+disp32] addressing
            if arch == Architecture::X86_64 && i + 2 < code.len() {
                let mut pos = i;
                let mut _rex_prefix = None;
                // Check if it's a REX prefix (0x40 - 0x4F)
                if code[pos] >= 0x40 && code[pos] <= 0x4F {
                    _rex_prefix = Some(code[pos]);
                    pos += 1;
                }

                if pos + 1 < code.len() {
                    let mut opcode = code[pos] as u32;
                    let mut modrm_pos = pos + 1;

                    // Handle 2-byte instructions (starting with 0x0F)
                    if opcode == 0x0F && pos + 2 < code.len() {
                        opcode = (opcode << 8) | (code[pos + 1] as u32);
                        modrm_pos = pos + 2;
                    }

                    if modrm_pos < code.len() {
                        let modrm = code[modrm_pos];

                        // Check if ModR/M is [RIP + disp32] (mod=00, rm=101)
                        if (modrm & 0xC7) == 0x05 && modrm_pos + 4 <= code.len() {
                            let disp_offset = modrm_pos + 1;
                            let next_instr_offset = disp_offset + 4;

                            // Branch 1: Indirect call/jmp (CALL/JMP [RIP+disp32]) -> point to Import Address Table (IAT)
                            // CALL: FF /2, JMP: FF /4
                            if opcode == 0xFF && ((modrm >> 3) & 7 == 2 || (modrm >> 3) & 7 == 4) {
                                // Read current displacement as import index
                                let import_index = u32::from_le_bytes([
                                    code[disp_offset],
                                    code[disp_offset + 1],
                                    code[disp_offset + 2],
                                    code[disp_offset + 3],
                                ]) as usize;

                                // Calculate actual IAT RVA
                                let mut current_iat_offset = 0;
                                let mut found_rva = None;
                                let mut flat_idx = 0;

                                'import_search: for (_, functions) in &self.imports {
                                    for _ in functions {
                                        if flat_idx == import_index {
                                            found_rva = Some(iat_start_rva + current_iat_offset);
                                            break 'import_search;
                                        }
                                        flat_idx += 1;
                                        current_iat_offset += pointer_size as u64;
                                    }
                                    current_iat_offset += pointer_size as u64; // skip DLL terminator
                                }

                                let target_rva =
                                    found_rva.unwrap_or(iat_start_rva + (import_index as u64 * pointer_size as u64));

                                let rip_rva: u64 = (code_section_rva as u64) + (next_instr_offset as u64);
                                let disp_i32 = (target_rva as i64 - rip_rva as i64) as i32;

                                code[disp_offset..disp_offset + 4].copy_from_slice(&disp_i32.to_le_bytes());

                                // Debug output
                                tracing::trace!(
                                    "IMPORT patch (RIP-relative): i={}, opcode={:04X}, rip_rva={:08X}, target_rva={:08X}, disp={:08X}, idx={}",
                                    i,
                                    opcode,
                                    rip_rva as u32,
                                    target_rva as u32,
                                    disp_i32 as u32,
                                    import_index
                                );

                                i = next_instr_offset;
                                continue;
                            }
                            // Branch 2: Data access instructions -> point to data section (.data)
                            // Any [RIP + disp32] that is not an import call is treated as data access
                            else {
                                // Special handling for LEA instruction relocation logic
                                // Check if it's LEA instruction (0x8D)
                                if opcode == 0x8D {
                                    // Check if REX prefix contains W bit (0x08)
                                    // For 64-bit LEA, REX should be 0x48 (W=1) or 0x4C (W=1, R=1) etc.
                                    let is_64_lea = if let Some(rex) = _rex_prefix { (rex & 0x08) != 0 } else { false };

                                    if is_64_lea {
                                        // Read current displacement as data index
                                        let data_index = u32::from_le_bytes([
                                            code[disp_offset],
                                            code[disp_offset + 1],
                                            code[disp_offset + 2],
                                            code[disp_offset + 3],
                                        ]) as usize;

                                        // Calculate data RVA
                                        let data_rva = data_section_rva.unwrap_or(code_section_rva + 0x1000) as u64;
                                        let target_rva = data_rva + (data_index as u64);
                                        let rip_rva: u64 = (code_section_rva as u64) + (next_instr_offset as u64);
                                        let disp_i32 = (target_rva as i64 - rip_rva as i64) as i32;

                                        code[disp_offset..disp_offset + 4].copy_from_slice(&disp_i32.to_le_bytes());

                                        // Debug output
                                        tracing::trace!(
                                            "LEA DATA patch (RIP-relative): i={}, opcode={:04X}, rip_rva={:08X}, target_rva={:08X}, disp={:08X}",
                                            i,
                                            opcode,
                                            rip_rva as u32,
                                            target_rva as u32,
                                            disp_i32 as u32
                                        );
                                    }
                                    else {
                                        // Non-64-bit LEA (e.g., using 32-bit register lea eax, [rip+disp32])
                                        // Read current displacement as offset within data section
                                        let current_disp = u32::from_le_bytes([
                                            code[disp_offset],
                                            code[disp_offset + 1],
                                            code[disp_offset + 2],
                                            code[disp_offset + 3],
                                        ]);

                                        let target_rva: u64 = (data_section_rva.unwrap_or(0) as u64) + (current_disp as u64);
                                        let rip_rva: u64 = (code_section_rva as u64) + (next_instr_offset as u64);
                                        let disp_i32 = (target_rva as i64 - rip_rva as i64) as i32;

                                        code[disp_offset..disp_offset + 4].copy_from_slice(&disp_i32.to_le_bytes());

                                        // Debug output
                                        tracing::trace!(
                                            "LEA DATA patch (RVA): i={}, opcode={:04X}, rip_rva={:08X}, target_rva={:08X}, disp={:08X}",
                                            i,
                                            opcode,
                                            rip_rva as u32,
                                            target_rva as u32,
                                            disp_i32 as u32
                                        );
                                    }
                                }
                                else {
                                    // Other [RIP + disp32] instructions (e.g., mov, cmp, etc.)
                                    // Read current displacement as offset within data section
                                    let current_disp = u32::from_le_bytes([
                                        code[disp_offset],
                                        code[disp_offset + 1],
                                        code[disp_offset + 2],
                                        code[disp_offset + 3],
                                    ]);

                                    let target_rva: u64 = (data_section_rva.unwrap_or(0) as u64) + (current_disp as u64);
                                    let rip_rva: u64 = (code_section_rva as u64) + (next_instr_offset as u64);
                                    let disp_i32 = (target_rva as i64 - rip_rva as i64) as i32;

                                    code[disp_offset..disp_offset + 4].copy_from_slice(&disp_i32.to_le_bytes());

                                    // Debug output
                                    tracing::trace!(
                                        "DATA patch (RVA): i={}, opcode={:04X}, rip_rva={:08X}, target_rva={:08X}, disp={:08X}",
                                        i,
                                        opcode,
                                        rip_rva as u32,
                                        target_rva as u32,
                                        disp_i32 as u32
                                    );
                                }
                            }

                            i = next_instr_offset;
                            continue;
                        }
                    }
                }
            }

            // x86 (32-bit) logic remains unchanged
            if arch == Architecture::X86 && i + 1 < code.len() && code[i] == 0xFF && code[i + 1] == 0x15 && i + 5 < code.len() {
                // Read current displacement as import index
                let import_index = u32::from_le_bytes([code[i + 2], code[i + 3], code[i + 4], code[i + 5]]) as usize;

                // Calculate actual IAT RVA
                let mut current_iat_offset = 0;
                let mut found_rva = None;
                let mut flat_idx = 0;

                'import_search_x86: for (_, functions) in &self.imports {
                    for _ in functions {
                        if flat_idx == import_index {
                            found_rva = Some(iat_start_rva + current_iat_offset);
                            break 'import_search_x86;
                        }
                        flat_idx += 1;
                        current_iat_offset += pointer_size as u64;
                    }
                    current_iat_offset += pointer_size as u64;
                }

                let target_rva = found_rva.unwrap_or(iat_start_rva + (import_index as u64 * pointer_size as u64));
                let target_va = image_base + target_rva;

                let disp: u32 = target_va as u32;
                let address_bytes = disp.to_le_bytes();
                code[i + 2..i + 6].copy_from_slice(&address_bytes);
                i += 6;
            }
            // Find direct CALL instruction (0xE8) - reserved for internal function calls
            else if code[i] == 0xE8 && i + 4 < code.len() {
                // Read relative offset
                let rel_offset = i32::from_le_bytes([code[i + 1], code[i + 2], code[i + 3], code[i + 4]]);

                // Debug output: Since AOT currently doesn't use internal labels, normal direct CALL shouldn't appear here unless it's a reserved placeholder.
                // If offset is 0, it might be an internal function label placeholder that needs patching.
                // We keep the rel_offset value for future expansion.
                tracing::trace!("Direct CALL (E8): i={}, rel_offset={:08X}", i, rel_offset);

                i += 5; // Skip the entire direct CALL instruction
            }
            // Find PUSH imm32 instruction (0x68) - used for push_label or data address placeholder
            else if code[i] == 0x68 && i + 4 < code.len() {
                // Read immediate value
                let imm = u32::from_le_bytes([code[i + 1], code[i + 2], code[i + 3], code[i + 4]]);

                if arch == Architecture::X86 {
                    // If this is a 0 placeholder, check if it should be patched to .data start address
                    if imm == 0 {
                        // Check if there's a push_imm8 + push_label pattern before
                        // push_imm(msg_len) compiles to 6a <len> (2 bytes)
                        // push_label("msg") compiles to 68 00 00 00 00 (5 bytes)
                        let mut should_patch = false;

                        // Look back to see if there's a push imm8 <msg_len> immediately followed by current push 0
                        if i >= 2 {
                            // Ensure enough space to look back
                            // Find previous push imm8 instruction (6a)
                            let prev_push_pos = i - 2;
                            if prev_push_pos < code.len() && code[prev_push_pos] == 0x6a {
                                let prev_imm8 = code[prev_push_pos + 1] as u32;

                                // Check if this immediate value equals message length
                                let msg_len = if let Some(data) = &self.data {
                                    data.iter().position(|&b| b == 0).map(|p| p as u32).unwrap_or(0)
                                }
                                else {
                                    0
                                };

                                if prev_imm8 == msg_len {
                                    should_patch = true;
                                }
                            }
                        }

                        if should_patch {
                            // Data section start VA
                            let data_section_va: u64 = image_base + (data_section_rva.unwrap_or(0) as u64);
                            let addr_u32 = data_section_va as u32;
                            code[i + 1..i + 5].copy_from_slice(&addr_u32.to_le_bytes());
                        }
                    }
                }

                i += 5; // Skip the entire PUSH instruction
            }
            else {
                i += 1;
            }
        }
    }

    /// Generate PE file byte array
    pub fn generate(&mut self) -> Result<Vec<u8>, GaiaError> {
        // Build sections
        self.sections = self.build_sections(); // Populate sections first

        // Build header
        let header = self.build_header()?;

        // Build import table
        let mut import_table = ImportTable::new();
        for (dll_name, functions) in &self.imports {
            let entry = ImportEntry { dll_name: dll_name.clone(), functions: functions.clone() };
            import_table.entries.push(entry);
        }

        // Create PE program
        let program = PeProgram {
            header,
            sections: self.sections.clone(),
            imports: import_table,
            exports: crate::types::tables::ExportTable::new(),
        };

        // Write to byte array
        let mut buffer = Vec::new();
        let cursor = Cursor::new(&mut buffer);
        let mut writer = ExeWriter::new(cursor);
        writer.write_program(&program)?;

        Ok(buffer)
    }
}

impl Default for PeBuilder {
    fn default() -> Self {
        Self::new()
    }
}