pe-assembler 0.1.1

use crate::types::{
    coff::{CoffFileType, CoffHeader, CoffInfo, CoffObject, CoffRelocation, CoffSymbol, SectionHeader},
    CoffSection,
};
use gaia_binary::{LittleEndian, ReadBytesExt};
use gaia_types::{helpers::Architecture, GaiaError};
use std::io::{Read, Seek, SeekFrom};

/// Common trait for COFF file readers
pub trait CoffReader<R: Read + Seek> {
    /// Get a mutable reference to the binary reader
    fn get_viewer(&mut self) -> &mut R;

    /// Get a mutable reference to diagnostics information
    fn add_diagnostics(&mut self, error: impl Into<GaiaError>);

    fn get_position(&mut self) -> Result<u64, GaiaError>
    where
        R: Seek,
    {
        Ok(self.get_viewer().stream_position()?)
    }

    fn set_position(&mut self, offset: u64) -> Result<u64, GaiaError>
    where
        R: Seek,
    {
        Ok(self.get_viewer().seek(SeekFrom::Start(offset))?)
    }

    /// Read COFF header info (generic implementation)
    fn get_coff_header(&mut self) -> Result<&CoffHeader, GaiaError>;

    /// Set COFF header info (generic implementation)
    fn set_coff_header(&mut self, head: CoffHeader) -> Option<CoffHeader>;

    /// Get cached section headers
    fn get_section_headers(&mut self) -> Result<&[SectionHeader], GaiaError>;

    /// Set cached section headers
    fn set_section_headers(&mut self, headers: Vec<SectionHeader>) -> Vec<SectionHeader>;

    /// Force read full COFF object and cache the result
    fn get_coff_object(&mut self) -> Result<&CoffObject, GaiaError>;

    /// Force read full COFF object and cache the result
    fn set_coff_object(&mut self, object: CoffObject) -> Option<CoffObject>;

    /// Get cached COFF info
    fn get_coff_info(&mut self) -> Result<&CoffInfo, GaiaError>;

    /// Set cached COFF info
    fn set_coff_info(&mut self, info: CoffInfo) -> Option<CoffInfo>;

    /// Create COFF info view (generic implementation)
    fn create_coff_info(&mut self) -> Result<CoffInfo, GaiaError> {
        let header = self.get_coff_header()?.clone();

        // Determine architecture based on machine type
        let target_arch = match header.machine {
            0x014c => Architecture::X86,
            0x8664 => Architecture::X86_64,
            0x01c0 => Architecture::ARM32,
            0xaa64 => Architecture::ARM64,
            _ => Architecture::Unknown,
        };

        // Get current file size
        let current_pos = self.get_position()?;
        self.get_viewer().seek(SeekFrom::End(0))?;
        let file_size = self.get_position()?;
        self.set_position(current_pos)?;

        Ok(CoffInfo {
            file_type: CoffFileType::Object,
            target_arch,
            section_count: header.number_of_sections,
            symbol_count: header.number_of_symbols,
            file_size,
            timestamp: header.time_date_stamp,
        })
    }
}

impl CoffRelocation {
    /// Read relocation entry from reader
    pub fn read<R: Read>(mut reader: R) -> Result<Self, GaiaError> {
        Ok(CoffRelocation {
            virtual_address: reader.read_u32::<LittleEndian>()?,
            symbol_table_index: reader.read_u32::<LittleEndian>()?,
            relocation_type: reader.read_u16::<LittleEndian>()?,
        })
    }
}

impl CoffSymbol {
    /// Read symbol from reader
    pub fn read<R: Read>(mut reader: R) -> Result<Self, GaiaError> {
        let mut name_bytes = [0u8; 8];
        reader.read_exact(&mut name_bytes)?;

        let name = if name_bytes[0..4] == [0, 0, 0, 0] {
            // Long name, stored in string table
            format!("symbol_{}", u32::from_le_bytes([name_bytes[4], name_bytes[5], name_bytes[6], name_bytes[7]]))
        }
        else {
            String::from_utf8_lossy(&name_bytes).trim_end_matches('\0').to_string()
        };

        Ok(CoffSymbol {
            name,
            value: reader.read_u32::<LittleEndian>()?,
            section_number: reader.read_i16::<LittleEndian>()?,
            symbol_type: reader.read_u16::<LittleEndian>()?,
            storage_class: reader.read_u8()?,
            number_of_aux_symbols: reader.read_u8()?,
        })
    }
}

impl SectionHeader {
    /// Read section header from reader
    pub fn read<R: Read>(mut reader: R) -> Result<Self, GaiaError> {
        let mut name = [0u8; 8];
        reader.read_exact(&mut name)?;

        Ok(SectionHeader {
            name,
            virtual_size: reader.read_u32::<LittleEndian>()?,
            virtual_address: reader.read_u32::<LittleEndian>()?,
            size_of_raw_data: reader.read_u32::<LittleEndian>()?,
            pointer_to_raw_data: reader.read_u32::<LittleEndian>()?,
            pointer_to_relocations: reader.read_u32::<LittleEndian>()?,
            pointer_to_line_numbers: reader.read_u32::<LittleEndian>()?,
            number_of_relocations: reader.read_u16::<LittleEndian>()?,
            number_of_line_numbers: reader.read_u16::<LittleEndian>()?,
            characteristics: reader.read_u32::<LittleEndian>()?,
        })
    }
}

/// Force read COFF header (generic implementation)
pub(crate) fn read_coff_header<R: Read + Seek>(reader: &mut impl CoffReader<R>) -> Result<CoffHeader, GaiaError> {
    // Save current position
    let original_pos = reader.get_position()?;

    // Reset to the beginning of the file
    reader.set_position(0)?;

    // Read COFF header
    let coff_header = CoffHeader::read(reader.get_viewer())?;

    // Verify machine type
    match coff_header.machine {
        0x014c | 0x8664 | 0x01c0 | 0xaa64 => {} // Supported architectures
        unknown => {
            let error = GaiaError::invalid_data(&format!("Unsupported machine type: 0x{:04x}", unknown));
            reader.add_diagnostics(error);
        }
    }

    // Verify section count
    if coff_header.number_of_sections == 0 {
        let error = GaiaError::invalid_data("COFF file must have at least one section");
        reader.add_diagnostics(error);
    }

    // Restore original position
    reader.set_position(original_pos)?;

    Ok(coff_header)
}

/// Read section header info (generic implementation)
pub(crate) fn read_section_headers<R: Read + Seek>(reader: &mut impl CoffReader<R>) -> Result<Vec<SectionHeader>, GaiaError> {
    // Read main header first
    let header = reader.get_coff_header()?.clone();
    let original_pos = reader.get_position()?;

    // Read section headers
    let mut section_headers = Vec::with_capacity(header.number_of_sections as usize);

    // Locate section header position (after COFF header)
    let section_header_offset = std::mem::size_of::<CoffHeader>() as u64 + header.size_of_optional_header as u64;

    reader.set_position(section_header_offset)?;

    for _ in 0..header.number_of_sections {
        let section_header = SectionHeader::read(reader.get_viewer())?;
        section_headers.push(section_header);
    }

    // Restore original position
    reader.set_position(original_pos)?;

    Ok(section_headers)
}

/// Read section data from section header (generic implementation)
pub(crate) fn read_section_from_header<R: Read + Seek>(
    reader: &mut impl CoffReader<R>,
    header: &SectionHeader,
) -> Result<CoffSection, GaiaError> {
    let mut data = Vec::new();
    let mut relocations = Vec::new();

    // Read section data
    if header.size_of_raw_data > 0 && header.pointer_to_raw_data > 0 {
        let original_pos = reader.get_position()?;
        reader.set_position(header.pointer_to_raw_data as u64)?;
        data.resize(header.size_of_raw_data as usize, 0);
        reader.get_viewer().read_exact(&mut data)?;
        reader.set_position(original_pos)?;
    }

    // Read relocation info
    if header.number_of_relocations > 0 && header.pointer_to_relocations > 0 {
        let original_pos = reader.get_position()?;
        reader.set_position(header.pointer_to_relocations as u64)?;

        for _ in 0..header.number_of_relocations {
            let relocation = CoffRelocation::read(reader.get_viewer())?;
            relocations.push(relocation);
        }

        reader.set_position(original_pos)?;
    }

    Ok(CoffSection { header: *header, data, relocations })
}

/// Read symbol table (generic implementation)
pub(crate) fn read_symbols<R: Read + Seek>(reader: &mut impl CoffReader<R>) -> Result<Vec<CoffSymbol>, GaiaError> {
    let header = reader.get_coff_header()?.clone();
    let original_pos = reader.get_position()?;

    // Calculate symbol table offset
    let symbol_table_offset = header.pointer_to_symbol_table as u64;
    reader.set_position(symbol_table_offset)?;

    let mut symbols = Vec::new();
    for _ in 0..header.number_of_symbols {
        let symbol = CoffSymbol::read(reader.get_viewer())?;
        symbols.push(symbol);
    }

    reader.set_position(original_pos)?;
    Ok(symbols)
}

/// Force read full COFF object
pub(crate) fn read_coff_object<R: Read + Seek>(reader: &mut impl CoffReader<R>) -> Result<CoffObject, GaiaError> {
    let header = reader.get_coff_header()?.clone();
    let section_headers = read_section_headers(reader)?;

    // Read section data
    let mut sections = Vec::new();
    for section_header in &section_headers {
        let section = read_section_from_header(reader, section_header)?;
        sections.push(section);
    }

    // Read symbol table
    let symbols = read_symbols(reader)?;

    // Read string table (simplified implementation)
    let string_table = Vec::new();

    Ok(CoffObject { header, sections, symbols, string_table })
}