use crate::app_error::AppError;
use crate::services::file::expand_file_size;
use memmap2::Mmap;
use pe_parse::header::{
    ImageDirectoryEntry, ImageImportDescriptor, ImageSectionHeader, ImportLookupTable32,
    ImportLookupTable64, SectionCharacteristics,
};
use pe_parse::pe::{MutablePE, ReadOnlyPE};
use serde::Serialize;
use crate::services::file;
use std::ffi::CStr;
use std::mem::{self};

/// 应用的状态结构体
#[derive(Default)]
pub struct AppState {
    /// 文件路径
    pub file_path: Option<String>,
    /// 文件Mmap的只读内存映射
    mmap: Option<Mmap>,
    /// 文件是否是64位
    pub is_64_bit: Option<bool>,
}

impl AppState {
    /// 初始化打开文件
    pub fn init_open_file(&mut self, file_path: &str) -> Result<(), AppError> {
        self.file_path = Some(file_path.to_string());
        self.mmap = Some(file::mmap_file(file_path)?);
        // 读取PE格式来判断是否是64位
        let mmap: &Mmap = self.mmap.as_ref().unwrap();
        self.is_64_bit = Some(mmap.is_64_bit()?);
        Ok(())
    }

    /// close_file 关闭文件
    pub fn close_file(&mut self) -> Result<(), AppError> {
        self.file_path = None;
        self.mmap = None;
        self.is_64_bit = None;
        Ok(())
    }

    /// 获取文件内存映射的引用
    pub fn get_mmap_ref(&self) -> Result<&Mmap, AppError> {
        self.mmap.as_ref().ok_or(AppError::NoFileOpened)
    }

    /// 设置文件内存映射
    pub fn set_mmap(&mut self, mmap: Mmap) {
        self.mmap = Some(mmap);
    }

    /// 在指定偏移处写入数据，此举会同步修改文件内容
    pub fn write_data(&mut self, offset: usize, data: &[u8]) -> Result<(), AppError> {
        // 1. 重新以可读可写打开文件
        let file_path = self.file_path.as_ref().ok_or(AppError::NoFileOpened)?;
        let mut rw_mmap = file::mmap_file_rw(file_path)?;
        // 2. 向内存映射中写入数据
        rw_mmap[offset..(offset + data.len())].copy_from_slice(data);
        rw_mmap.flush()?;

        // 3. 重新打开文件
        self.mmap = Some(file::mmap_file(file_path)?);
        self.is_64_bit = Some(self.get_mmap_ref()?.is_64_bit()?);

        Ok(())
    }

    pub fn add_section(
        &mut self,
        section_name: &str,
        section_size: u32,
        section_characteristics: u32,
    ) -> Result<(), AppError> {
        if section_name.len() > 8 {
            // section_name不能超过8个字符
            return Err(AppError::SectionNameTooLong);
        }
        let mut section_name_bytes = [0u8; 8];
        section_name_bytes[..section_name.len()].copy_from_slice(section_name.as_bytes());
        let mmap = self.get_mmap_ref()?;
        let file_path = self.file_path.as_ref().ok_or(AppError::NoFileOpened)?;
        // 将section_size 进行文件对齐
        let file_alignment = mmap.align_size_with_file_alignment(section_size)?;
        // 判断是否可以添加节
        // 1. 判断头部空余空间是否足够
        const SIZE_OF_SECTION_HEADER: u32 = mem::size_of::<ImageSectionHeader>() as u32;
        let header_empty_space = mmap.get_header_empty_space_size()?;
        let mut rw_mmap_option = None;

        if header_empty_space < SIZE_OF_SECTION_HEADER {
            // 头部空余空间不足 需要判断是否可以拓展
            // 获取最大可拓展的大小
            let max_expand_size = mmap.get_header_max_expand_size()?;
            let file_alignment = mmap.align_size_with_file_alignment(SIZE_OF_SECTION_HEADER)?;
            if file_alignment > max_expand_size {
                return Err(AppError::CannotModifyFile("无法拓展文件大小!".to_string()));
            }
            // 否则可以通过拓展头部大小的方式进行节的添加
            // 1. 拓展文件大小
            expand_file_size(file_path, file_alignment)?;

            // 2. 以可读可写的方式重新映射文件
            rw_mmap_option = Some(file::mmap_file_rw(file_path)?);
            // 拓展头部大小 拓展一个节的大小
            rw_mmap_option
                .as_mut()
                .unwrap()
                .expand_headers(file_alignment)?;
        }
        // 2. 添加节
        // 先拓展文件大小
        expand_file_size(file_path, file_alignment)?;
        // 如果rw_mmap_option为None，则需要重新映射文件
        if rw_mmap_option.is_none() {
            rw_mmap_option = Some(file::mmap_file_rw(self.file_path.as_ref().unwrap())?);
        }
        let rw_mmap = rw_mmap_option.as_mut().unwrap();

        let characteristics = SectionCharacteristics::from_bits(section_characteristics).unwrap();

        rw_mmap.add_section(&section_name_bytes, section_size, characteristics)?;

        // 3. 重新映射文件
        self.mmap = Some(file::mmap_file(self.file_path.as_ref().unwrap())?);

        Ok(())
    }

    pub fn get_import_tables(&self) -> Result<Vec<ImportModuleTable>, AppError> {
        let mmap = self.get_mmap_ref()?;
        let is_64_bit = self.is_64_bit.unwrap();
        // 1. 获取导入表
        let import_table = mmap.get_data_directory(ImageDirectoryEntry::Import)?;
        // 2. 获取导入表的RVA
        let import_table_rva = import_table.virtual_address;
        // 3. RVA转FOA
        let import_table_foa = mmap.rva_to_foa(import_table_rva)?;
        // 从文件import_table_foa循环解析ImageImportDescriptor，当解析出的导入表==ImageImportDescriptor::default()时，结束循环
        let mut origin_import_module_tables: Vec<&ImageImportDescriptor> = Vec::new();
        let mut base_offset = import_table_foa as usize;
        loop {
            let import_module_item = unsafe {
                let ptr = mmap.as_ptr().add(base_offset) as *const ImageImportDescriptor;
                ptr.as_ref().unwrap()
            };
            if import_module_item == &ImageImportDescriptor::default() {
                break;
            }
            origin_import_module_tables.push(import_module_item);
            base_offset += mem::size_of::<ImageImportDescriptor>();
        }
        // 4. 解析导入表 主要是来格式化一些数据 先创建一个ImportModuleTable 的变长数组
        let mut import_module_tables: Vec<ImportModuleTable> = Vec::new();
        // 循环处理origin_import_module_tables
        for origin_item in origin_import_module_tables.iter() {
            let mut import_module = ImportModuleTable::default();
            // 1. 获取模块名称
            // 先获取模块名称的FOA
            let module_name_foa = mmap.rva_to_foa(origin_item.name)?;
            // 获取模块名称 需要unsafe读取
            let module_name = unsafe {
                let ptr = mmap[module_name_foa as usize..].as_ptr() as *const i8;
                CStr::from_ptr(ptr).to_str()?.to_string()
            };
            import_module.module_name = module_name;
            // 其他成员
            import_module.forwarder_chain = origin_item.forwarder_chain;
            import_module.iat_rva = origin_item.first_thunk;
            import_module.int_rva = origin_item.original_first_thunk;
            import_module.module_name_rva = origin_item.name;
            import_module.timestamp = origin_item.time_date_stamp;
            // 2. 获取函数表
            let mut functions: Vec<ImportFunctionTable> = Vec::new();
            // 获取函数表的FOA
            let mut base_offset = mmap.rva_to_foa(origin_item.original_first_thunk)? as usize;
            // 循环解析函数表
            loop {
                // TODO: 这里需要判断是64位还是32位
                // 如果是64位，每次读取8字节 否则读取4字节
                let import_lookup_table = if is_64_bit {
                    let ptr = mmap[base_offset as usize..].as_ptr() as *const u64;
                    unsafe { ptr.read() as u64 }
                } else {
                    let ptr = mmap[base_offset as usize..].as_ptr() as *const u32;
                    unsafe { ptr.read() as u64 }
                };

                // 如果import_lookup_table == 0 则表示结束
                if import_lookup_table == 0 {
                    break;
                }

                // 1. 判断是序号还是命名函数
                let is_named = match is_64_bit {
                    true => {
                        import_lookup_table & ImportLookupTable64::IMPORT_BY_ORDINAL.bits() == 0
                    }
                    false => {
                        import_lookup_table as u32 & ImportLookupTable32::IMPORT_BY_ORDINAL.bits()
                            == 0
                    }
                };
                let mut import_function_table = ImportFunctionTable::default();

                if is_named == true {
                    // TODO: 仅当序号/名称标志位字段为 0（按名称导入）时，才使用此字段。
                    let hint_name_table_rva = import_lookup_table as u32 & 0x7FFFFFFF;
                    let hint_name_table_foa = mmap.rva_to_foa(hint_name_table_rva)?;
                    // 1. 读hint 两个字节
                    let funtion_hint: u16 = unsafe {
                        let ptr = mmap[hint_name_table_foa as usize..].as_ptr() as *const u16;
                        ptr.read()
                    };
                    import_function_table.function_hint = funtion_hint;
                    import_function_table.function_address = hint_name_table_rva;
                    let function_name = unsafe {
                        let ptr = mmap[hint_name_table_foa as usize + 2..].as_ptr() as *const i8;
                        CStr::from_ptr(ptr).to_str()?.to_string()
                    };
                    import_function_table.function_name = function_name;
                } else {
                    let function_ordianl = import_lookup_table as u32
                        & ImportLookupTable32::IMPORT_BY_ORDINAL_MASK.bits();
                    import_function_table.function_hint = function_ordianl as u16;
                    import_function_table.function_type = ImportFunctionTableItem::Ordinal;
                }
                functions.push(import_function_table);
                // 下一个
                base_offset += if is_64_bit { 8 } else { 4 };
            }
            import_module.functions = functions;
            import_module_tables.push(import_module);
        }
        Ok(import_module_tables)
    }
}

// 定义一个获取导入表后需要返回的结构
#[derive(Debug, Default, Clone, Serialize)]
pub struct ImportModuleTable {
    pub module_name: String,                 // 模块名称
    pub functions: Vec<ImportFunctionTable>, // 这里要另一个结构体来描述
    pub timestamp: u32,                      // 时间戳
    pub forwarder_chain: u32,                // 转发链
    pub module_name_rva: u32,                // dll名称的RVA
    pub iat_rva: u32,                        // IAT的RVA
    pub int_rva: u32,                        // INT的RVA
}

// 因为导入查找表可能是函数名、也可能是序号 所以用枚举包一下最好
#[derive(Debug, Clone, Serialize, Default)]
pub enum ImportFunctionTableItem {
    // 命名的导入函数
    #[default]
    Named,
    // 序号表示的
    Ordinal,
}

#[derive(Debug, Default, Clone, Serialize)]
pub struct ImportFunctionTable {
    pub function_name: String, // 函数名称
    pub function_address: u32, // 函数地址 IAT
    pub function_hint: u16,    // 函数提示
    pub function_type: ImportFunctionTableItem,
}
