//! COFF Object File Parser
//!
//! 解析 MSVC cl.exe / clang-cl 生成的 COFF .obj 文件
//! 支持读取节区头、符号表、字符串表、重定位项

use anyhow::{Result, bail};

/// COFF 文件头 (20 bytes)
#[derive(Debug, Clone)]
pub struct CoffHeader {
    pub machine: u16,
    pub num_sections: u16,
    pub time_date_stamp: u32,
    pub ptr_symbol_table: u32,
    pub num_symbols: u32,
    pub size_optional_header: u16,
    pub characteristics: u16,
}

/// 节区头 (40 bytes)
#[derive(Debug, Clone)]
pub struct SectionHeader {
    pub name: [u8; 8],
    pub virtual_size: u32,
    pub virtual_address: u32,
    pub size_of_raw_data: u32,
    pub ptr_to_raw_data: u32,
    pub ptr_to_relocations: u32,
    pub ptr_to_line_numbers: u32,
    pub num_relocations: u16,
    pub num_line_numbers: u16,
    pub characteristics: u32,
}

/// COFF 符号表项 (18 bytes)
#[derive(Debug, Clone)]
pub struct Symbol {
    pub name: String,
    pub value: u32,
    pub section_number: i16,
    pub sym_type: u16,
    pub storage_class: u8,
    pub num_aux_symbols: u8,
    pub coff_index: u32, // 在 COFF 文件中的原始符号索引（包括辅助记录）
}

/// 重定位项 (10 bytes)
#[derive(Debug, Clone)]
pub struct Relocation {
    pub virtual_address: u32,
    pub symbol_table_index: u32,
    pub rel_type: u16,
}

/// 解析后的 COFF 对象文件
#[derive(Debug, Clone)]
pub struct CoffObject {
    pub header: CoffHeader,
    pub sections: Vec<Section>,
    pub symbols: Vec<Symbol>,
    pub string_table: Vec<u8>,
}

/// 节区（含原始数据）
#[derive(Debug, Clone)]
pub struct Section {
    pub header: SectionHeader,
    pub data: Vec<u8>,
    pub relocations: Vec<Relocation>,
}

impl CoffObject {
    /// 从字节数组解析 COFF 对象文件
    pub fn parse(data: &[u8]) -> Result<Self> {
        if data.len() < 20 {
            bail!("COFF file too small ({} bytes)", data.len());
        }

        // 解析文件头
        let header = CoffHeader {
            machine: read_u16(data, 0),
            num_sections: read_u16(data, 2),
            time_date_stamp: read_u32(data, 4),
            ptr_symbol_table: read_u32(data, 8),
            num_symbols: read_u32(data, 12),
            size_optional_header: read_u16(data, 16),
            characteristics: read_u16(data, 18),
        };

        // 只支持 AMD64
        if header.machine != 0x8664 {
            bail!("Unsupported machine type: 0x{:04X} (expected AMD64=0x8664)", header.machine);
        }

        let mut offset = 20 + header.size_optional_header as usize;

        // 解析节区头
        let mut sections = Vec::new();
        for _ in 0..header.num_sections {
            if offset + 40 > data.len() {
                bail!("Section header extends beyond file");
            }

            let section_header = SectionHeader {
                name: {
                    let mut name = [0u8; 8];
                    name.copy_from_slice(&data[offset..offset + 8]);
                    name
                },
                virtual_size: read_u32(data, offset + 8),
                virtual_address: read_u32(data, offset + 12),
                size_of_raw_data: read_u32(data, offset + 16),
                ptr_to_raw_data: read_u32(data, offset + 20),
                ptr_to_relocations: read_u32(data, offset + 24),
                ptr_to_line_numbers: read_u32(data, offset + 28),
                num_relocations: read_u16(data, offset + 32),
                num_line_numbers: read_u16(data, offset + 34),
                characteristics: read_u32(data, offset + 36),
            };

            // 读取节区数据
            let section_data = if section_header.ptr_to_raw_data > 0 && section_header.size_of_raw_data > 0 {
                let start = section_header.ptr_to_raw_data as usize;
                let end = start + section_header.size_of_raw_data as usize;
                if end > data.len() {
                    bail!("Section data extends beyond file");
                }
                data[start..end].to_vec()
            } else {
                Vec::new()
            };

            // 读取重定位项
            let mut relocations = Vec::new();
            if section_header.ptr_to_relocations > 0 && section_header.num_relocations > 0 {
                let mut rel_offset = section_header.ptr_to_relocations as usize;
                for _ in 0..section_header.num_relocations {
                    if rel_offset + 10 > data.len() {
                        bail!("Relocation extends beyond file");
                    }
                    relocations.push(Relocation {
                        virtual_address: read_u32(data, rel_offset),
                        symbol_table_index: read_u32(data, rel_offset + 4),
                        rel_type: read_u16(data, rel_offset + 8),
                    });
                    rel_offset += 10;
                }
            }

            sections.push(Section {
                header: section_header,
                data: section_data,
                relocations,
            });

            offset += 40;
        }

        // 读取字符串表
        let string_table_offset = header.ptr_symbol_table as usize + (header.num_symbols as usize * 18);
        let string_table = if string_table_offset + 4 <= data.len() {
            let str_table_size = read_u32(data, string_table_offset) as usize;
            if string_table_offset + str_table_size <= data.len() {
                data[string_table_offset..string_table_offset + str_table_size].to_vec()
            } else {
                Vec::new()
            }
        } else {
            Vec::new()
        };

        // 解析符号表
        let mut symbols = Vec::new();
        let mut sym_offset = header.ptr_symbol_table as usize;
        let mut i = 0;
        while i < header.num_symbols {
            if sym_offset + 18 > data.len() {
                bail!("Symbol extends beyond file");
            }

            let name = read_symbol_name(data, sym_offset, &string_table)?;
            let value = read_u32(data, sym_offset + 8);
            let section_number = read_i16(data, sym_offset + 12);
            let sym_type = read_u16(data, sym_offset + 14);
            let storage_class = data[sym_offset + 16];
            let num_aux = data[sym_offset + 17];

            symbols.push(Symbol {
                name,
                value,
                section_number,
                sym_type,
                storage_class,
                num_aux_symbols: num_aux,
                coff_index: i, // 当前在 COFF 文件中的原始索引
            });

            sym_offset += 18 * (1 + num_aux as usize);
            i += 1 + num_aux as u32;
        }

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

    /// 获取所有外部符号（需要其他对象文件或库提供的）
    pub fn undefined_symbols(&self) -> Vec<&Symbol> {
        self.symbols.iter()
            .filter(|s| s.storage_class == 2 && s.section_number == 0) // EXTERNAL + UNDEF
            .collect()
    }

    /// 获取所有已定义的全局符号
    pub fn defined_symbols(&self) -> Vec<&Symbol> {
        self.symbols.iter()
            .filter(|s| {
                (s.storage_class == 2 || s.storage_class == 3) && // EXTERNAL or STATIC
                s.section_number > 0
            })
            .collect()
    }

    /// 按名称查找符号
    pub fn find_symbol(&self, name: &str) -> Option<&Symbol> {
        self.symbols.iter().find(|s| s.name == name)
    }
}

/// 读取符号名称（处理短名和长名）
fn read_symbol_name(data: &[u8], offset: usize, string_table: &[u8]) -> Result<String> {
    let first_four = read_u32(data, offset);
    let second_four = read_u32(data, offset + 4);

    if first_four == 0 {
        // 长名：偏移量指向字符串表
        let str_offset = second_four as usize;
        if str_offset < string_table.len() {
            let mut end = str_offset;
            while end < string_table.len() && string_table[end] != 0 {
                end += 1;
            }
            return Ok(String::from_utf8_lossy(&string_table[str_offset..end]).to_string());
        } else {
            bail!("String table offset out of bounds: {}", str_offset);
        }
    } else {
        // 短名：直接读取8字节
        let name_bytes = &data[offset..offset + 8];
        let mut len = 8;
        while len > 0 && name_bytes[len - 1] == 0 {
            len -= 1;
        }
        Ok(String::from_utf8_lossy(&name_bytes[..len]).to_string())
    }
}

fn read_u16(data: &[u8], offset: usize) -> u16 {
    u16::from_le_bytes([data[offset], data[offset + 1]])
}

fn read_u32(data: &[u8], offset: usize) -> u32 {
    u32::from_le_bytes([data[offset], data[offset + 1], data[offset + 2], data[offset + 3]])
}

fn read_i16(data: &[u8], offset: usize) -> i16 {
    i16::from_le_bytes([data[offset], data[offset + 1]])
}

/// 节区特征标志
pub mod section_flags {
    pub const IMAGE_SCN_CNT_CODE: u32 = 0x00000020;
    pub const IMAGE_SCN_CNT_INITIALIZED_DATA: u32 = 0x00000040;
    pub const IMAGE_SCN_CNT_UNINITIALIZED_DATA: u32 = 0x00000080;
    pub const IMAGE_SCN_MEM_EXECUTE: u32 = 0x20000000;
    pub const IMAGE_SCN_MEM_READ: u32 = 0x40000000;
    pub const IMAGE_SCN_MEM_WRITE: u32 = 0x80000000;
}

/// x64 重定位类型
pub mod reloc_types {
    pub const IMAGE_REL_AMD64_ABSOLUTE: u16 = 0x0000;
    pub const IMAGE_REL_AMD64_ADDR64: u16 = 0x0001;
    pub const IMAGE_REL_AMD64_ADDR32: u16 = 0x0002;
    pub const IMAGE_REL_AMD64_ADDR32NB: u16 = 0x0003;
    pub const IMAGE_REL_AMD64_REL32: u16 = 0x0004;
    pub const IMAGE_REL_AMD64_REL32_1: u16 = 0x0005;
    pub const IMAGE_REL_AMD64_REL32_2: u16 = 0x0006;
    pub const IMAGE_REL_AMD64_REL32_3: u16 = 0x0007;
    pub const IMAGE_REL_AMD64_REL32_4: u16 = 0x0008;
    pub const IMAGE_REL_AMD64_REL32_5: u16 = 0x0009;
    pub const IMAGE_REL_AMD64_SECTION: u16 = 0x000A;
    pub const IMAGE_REL_AMD64_SECREL: u16 = 0x000B;
}

/// 符号存储类
pub mod storage_class {
    pub const IMAGE_SYM_CLASS_END_OF_FUNCTION: u8 = 0xFF;
    pub const IMAGE_SYM_CLASS_NULL: u8 = 0;
    pub const IMAGE_SYM_CLASS_AUTOMATIC: u8 = 1;
    pub const IMAGE_SYM_CLASS_EXTERNAL: u8 = 2;
    pub const IMAGE_SYM_CLASS_STATIC: u8 = 3;
    pub const IMAGE_SYM_CLASS_REGISTER: u8 = 4;
    pub const IMAGE_SYM_CLASS_EXTERNAL_DEF: u8 = 5;
    pub const IMAGE_SYM_CLASS_LABEL: u8 = 6;
    pub const IMAGE_SYM_CLASS_UNDEFINED_LABEL: u8 = 7;
    pub const IMAGE_SYM_CLASS_MEMBER_OF_STRUCT: u8 = 8;
    pub const IMAGE_SYM_CLASS_ARGUMENT: u8 = 9;
    pub const IMAGE_SYM_CLASS_STRUCT_TAG: u8 = 10;
    pub const IMAGE_SYM_CLASS_MEMBER_OF_UNION: u8 = 11;
    pub const IMAGE_SYM_CLASS_UNION_TAG: u8 = 12;
    pub const IMAGE_SYM_CLASS_TYPE_DEFINITION: u8 = 13;
    pub const IMAGE_SYM_CLASS_UNDEFINED_STATIC: u8 = 14;
    pub const IMAGE_SYM_CLASS_ENUM_TAG: u8 = 15;
    pub const IMAGE_SYM_CLASS_MEMBER_OF_ENUM: u8 = 16;
    pub const IMAGE_SYM_CLASS_REGISTER_PARAM: u8 = 17;
    pub const IMAGE_SYM_CLASS_BIT_FIELD: u8 = 18;
    pub const IMAGE_SYM_CLASS_BLOCK: u8 = 100;
    pub const IMAGE_SYM_CLASS_FUNCTION: u8 = 101;
    pub const IMAGE_SYM_CLASS_END_OF_STRUCT: u8 = 102;
    pub const IMAGE_SYM_CLASS_FILE: u8 = 103;
    pub const IMAGE_SYM_CLASS_SECTION: u8 = 104;
    pub const IMAGE_SYM_CLASS_WEAK_EXTERNAL: u8 = 105;
    pub const IMAGE_SYM_CLASS_CLR_TOKEN: u8 = 107;
}

#[cfg(test)]
mod tests {
    use super::*;

    // === 辅助函数：构造最小有效 COFF 数据 ===

    /// 构造包含 1 个 text 节区 + 1 个符号的最小 COFF
    fn make_minimal_coff(section_name: &str, section_data: &[u8], symbol_name: &str) -> Vec<u8> {
        let mut data = Vec::new();

        // COFF Header (20 bytes)
        data.extend_from_slice(&0x8664u16.to_le_bytes()); // machine: AMD64
        data.extend_from_slice(&1u16.to_le_bytes());     // num_sections: 1
        data.extend_from_slice(&0u32.to_le_bytes());     // time_date_stamp
        let sym_table_offset = 20u32 + 40 + section_data.len() as u32; // header + section header + section data
        data.extend_from_slice(&sym_table_offset.to_le_bytes()); // ptr_symbol_table
        data.extend_from_slice(&1u32.to_le_bytes());     // num_symbols: 1
        data.extend_from_slice(&0u16.to_le_bytes());     // size_optional_header
        data.extend_from_slice(&0u16.to_le_bytes());     // characteristics

        // Section Header (40 bytes)
        let mut name = [0u8; 8];
        let name_bytes = section_name.as_bytes();
        let len = name_bytes.len().min(8);
        name[..len].copy_from_slice(&name_bytes[..len]);
        data.extend_from_slice(&name);
        data.extend_from_slice(&(section_data.len() as u32).to_le_bytes()); // virtual_size
        data.extend_from_slice(&0u32.to_le_bytes());     // virtual_address
        data.extend_from_slice(&(section_data.len() as u32).to_le_bytes()); // size_of_raw_data
        data.extend_from_slice(&(20u32 + 40u32).to_le_bytes()); // ptr_to_raw_data
        data.extend_from_slice(&0u32.to_le_bytes());     // ptr_to_relocations
        data.extend_from_slice(&0u32.to_le_bytes());     // ptr_to_line_numbers
        data.extend_from_slice(&0u16.to_le_bytes());     // num_relocations
        data.extend_from_slice(&0u16.to_le_bytes());     // num_line_numbers
        let chars = section_flags::IMAGE_SCN_CNT_CODE | section_flags::IMAGE_SCN_MEM_EXECUTE | section_flags::IMAGE_SCN_MEM_READ;
        data.extend_from_slice(&chars.to_le_bytes()); // characteristics

        // Section Data
        data.extend_from_slice(section_data);

        // Symbol Table (18 bytes)
        let mut sym_name = [0u8; 8];
        let sym_bytes = symbol_name.as_bytes();
        let sym_len = sym_bytes.len().min(8);
        sym_name[..sym_len].copy_from_slice(&sym_bytes[..sym_len]);
        data.extend_from_slice(&sym_name);          // name (8 bytes)
        data.extend_from_slice(&0u32.to_le_bytes()); // value
        data.extend_from_slice(&1i16.to_le_bytes()); // section_number: 1 (first section)
        data.extend_from_slice(&0u16.to_le_bytes()); // sym_type
        data.push(storage_class::IMAGE_SYM_CLASS_EXTERNAL); // storage_class
        data.push(0); // num_aux_symbols

        // String Table (4 bytes minimum - size field)
        data.extend_from_slice(&4u32.to_le_bytes());

        data
    }

    // === 基础解析测试 ===

    #[test]
    fn test_parse_too_small() {
        let data = vec![0u8; 19]; // 小于 20 字节
        assert!(CoffObject::parse(&data).is_err());
    }

    #[test]
    fn test_parse_unsupported_machine() {
        let mut data = vec![0u8; 20];
        data[0] = 0x01; data[1] = 0x00; // i386, not AMD64
        assert!(CoffObject::parse(&data).is_err());
    }

    #[test]
    fn test_parse_minimal_valid() {
        let data = make_minimal_coff(".text", &[0x90, 0xC3], "main");
        let coff = CoffObject::parse(&data).unwrap();
        assert_eq!(coff.header.machine, 0x8664);
        assert_eq!(coff.sections.len(), 1);
        assert_eq!(coff.symbols.len(), 1);
    }

    #[test]
    fn test_parse_section_name() {
        let data = make_minimal_coff(".text", &[0x90], "main");
        let coff = CoffObject::parse(&data).unwrap();
        let name = String::from_utf8_lossy(&coff.sections[0].header.name)
            .trim_end_matches('\0').to_string();
        assert_eq!(name, ".text");
    }

    #[test]
    fn test_parse_section_data() {
        let code = vec![0x90, 0x90, 0xC3];
        let data = make_minimal_coff(".text", &code, "main");
        let coff = CoffObject::parse(&data).unwrap();
        assert_eq!(coff.sections[0].data, code);
    }

    // === 符号表测试 ===

    #[test]
    fn test_parse_short_symbol_name() {
        let data = make_minimal_coff(".text", &[0xC3], "foo");
        let coff = CoffObject::parse(&data).unwrap();
        assert_eq!(coff.symbols[0].name, "foo");
        assert_eq!(coff.symbols[0].section_number, 1);
        assert_eq!(coff.symbols[0].storage_class, storage_class::IMAGE_SYM_CLASS_EXTERNAL);
    }

    #[test]
    fn test_find_symbol() {
        let data = make_minimal_coff(".text", &[0xC3], "my_func");
        let coff = CoffObject::parse(&data).unwrap();
        assert!(coff.find_symbol("my_func").is_some());
        assert!(coff.find_symbol("nonexistent").is_none());
    }

    #[test]
    fn test_defined_symbols() {
        let data = make_minimal_coff(".text", &[0xC3], "defined"); // 7 chars, fits in short name
        let coff = CoffObject::parse(&data).unwrap();
        let defined = coff.defined_symbols();
        assert_eq!(defined.len(), 1);
        assert_eq!(defined[0].name, "defined");
    }

    #[test]
    fn test_undefined_symbols() {
        // 构造一个未定义符号（section_number = 0, storage_class = EXTERNAL）
        let mut data = Vec::new();
        // Header
        data.extend_from_slice(&0x8664u16.to_le_bytes());
        data.extend_from_slice(&0u16.to_le_bytes()); // 0 sections
        data.extend_from_slice(&0u32.to_le_bytes());
        data.extend_from_slice(&20u32.to_le_bytes()); // ptr_symbol_table (right after header)
        data.extend_from_slice(&1u32.to_le_bytes()); // 1 symbol
        data.extend_from_slice(&0u16.to_le_bytes());
        data.extend_from_slice(&0u16.to_le_bytes());
        // Symbol: undefined external
        let mut sym_name = [0u8; 8];
        sym_name[..3].copy_from_slice(b"ext");
        data.extend_from_slice(&sym_name);
        data.extend_from_slice(&0u32.to_le_bytes()); // value
        data.extend_from_slice(&0i16.to_le_bytes()); // section_number: 0 = UNDEF
        data.extend_from_slice(&0u16.to_le_bytes()); // sym_type
        data.push(storage_class::IMAGE_SYM_CLASS_EXTERNAL);
        data.push(0);
        // String table
        data.extend_from_slice(&4u32.to_le_bytes());

        let coff = CoffObject::parse(&data).unwrap();
        let undefined = coff.undefined_symbols();
        assert_eq!(undefined.len(), 1);
        assert_eq!(undefined[0].name, "ext");
    }

    // === 长符号名测试 ===

    #[test]
    fn test_parse_long_symbol_name() {
        let long_name = "very_long_function_name_that_exceeds_8_bytes";
        let mut data = Vec::new();
        // Header
        data.extend_from_slice(&0x8664u16.to_le_bytes());
        data.extend_from_slice(&0u16.to_le_bytes()); // 0 sections
        data.extend_from_slice(&0u32.to_le_bytes());
        let sym_offset = 20u32;
        data.extend_from_slice(&sym_offset.to_le_bytes());
        data.extend_from_slice(&1u32.to_le_bytes()); // 1 symbol
        data.extend_from_slice(&0u16.to_le_bytes());
        data.extend_from_slice(&0u16.to_le_bytes());
        // Symbol with long name (first 4 bytes = 0, next 4 = string table offset)
        data.extend_from_slice(&[0u8; 4]); // zero first four = long name
        data.extend_from_slice(&4u32.to_le_bytes()); // offset 4 in string table
        data.extend_from_slice(&0u32.to_le_bytes()); // value
        data.extend_from_slice(&0i16.to_le_bytes()); // section_number
        data.extend_from_slice(&0u16.to_le_bytes());
        data.push(storage_class::IMAGE_SYM_CLASS_EXTERNAL);
        data.push(0);
        // String table: size + string (starting at offset 4) + null terminator
        let str_content = format!("{}\0", long_name);
        let str_table_size = 4 + str_content.len() as u32;
        data.extend_from_slice(&str_table_size.to_le_bytes());
        data.extend_from_slice(str_content.as_bytes());

        let coff = CoffObject::parse(&data).unwrap();
        assert_eq!(coff.symbols[0].name, long_name);
    }

    // === read 辅助函数测试 ===

    #[test]
    fn test_read_u16() {
        let data = [0x34, 0x12];
        assert_eq!(read_u16(&data, 0), 0x1234);
    }

    #[test]
    fn test_read_u32() {
        let data = [0x78, 0x56, 0x34, 0x12];
        assert_eq!(read_u32(&data, 0), 0x12345678);
    }

    #[test]
    fn test_read_i16() {
        let data = [0x00, 0x80]; // -32768
        assert_eq!(read_i16(&data, 0), -32768i16);
    }

    // === 常量测试 ===

    #[test]
    fn test_section_flags() {
        assert_eq!(section_flags::IMAGE_SCN_CNT_CODE, 0x00000020);
        assert_eq!(section_flags::IMAGE_SCN_CNT_INITIALIZED_DATA, 0x00000040);
        assert_eq!(section_flags::IMAGE_SCN_MEM_EXECUTE, 0x20000000);
        assert_eq!(section_flags::IMAGE_SCN_MEM_READ, 0x40000000);
        assert_eq!(section_flags::IMAGE_SCN_MEM_WRITE, 0x80000000);
    }

    #[test]
    fn test_reloc_types() {
        assert_eq!(reloc_types::IMAGE_REL_AMD64_ADDR64, 0x0001);
        assert_eq!(reloc_types::IMAGE_REL_AMD64_REL32, 0x0004);
        assert_eq!(reloc_types::IMAGE_REL_AMD64_SECTION, 0x000A);
    }

    #[test]
    fn test_storage_class() {
        assert_eq!(storage_class::IMAGE_SYM_CLASS_NULL, 0);
        assert_eq!(storage_class::IMAGE_SYM_CLASS_EXTERNAL, 2);
        assert_eq!(storage_class::IMAGE_SYM_CLASS_STATIC, 3);
        assert_eq!(storage_class::IMAGE_SYM_CLASS_FILE, 103);
    }
}
