use std::io::Seek;
use std::option;

use crate::app_error::AppError;
use crate::{pe_parse::pe::ReadOnlyPE, services::file};
use memmap2::{Mmap, MmapOptions};

/// 应用的状态结构体
#[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: usize,
        section_characteristics: u32,
    ) -> Result<(), AppError> {
        // 1. section_name不能超过8个字符
        if section_name.len() > 8 {
            return Err(AppError::SectionNameTooLong);
        }
        // as_bytes 没有补0
        let mut section_name_bytes = [0u8; 8];
        // 循环拷贝section_name的数据
        for i in 0..section_name.len() {
            if i < section_name.len() {
                section_name_bytes[i] = section_name.as_bytes()[i];
            }
        }

        // 判断是否可以添加节
        // if(size_of_headers - 最后一个节表的偏移 - 40) >= 40，则可以添加节表。
        // 2. 获取size_of_headers
        let mmap = self.get_mmap_ref()?;
        let origin_file_len = mmap.len();
        let optional_header = mmap.get_optional_header()?;
        let file_header_offset = mmap.get_file_header_offset()?;
        let optional_header_offset = mmap.get_optional_header_offset()?;
        let file_header = mmap.get_file_header()?;
        let size_of_headers = match optional_header {
            crate::pe_parse::header::ImageOptionalHeader::OptionalHeader32(
                image_optional_header32,
            ) => image_optional_header32.size_of_headers,
            crate::pe_parse::header::ImageOptionalHeader::OptionalHeader64(
                image_optional_header64,
            ) => image_optional_header64.size_of_headers,
        };
        let section_table_offset = mmap.get_section_headers_offset()?;
        // 获取节表的数量
        let number_of_sections = file_header.number_of_sections;
        // 1. 获取偏移
        let section_alignment = match optional_header {
            crate::pe_parse::header::ImageOptionalHeader::OptionalHeader32(
                image_optional_header32,
            ) => image_optional_header32.section_alignment,
            crate::pe_parse::header::ImageOptionalHeader::OptionalHeader64(
                image_optional_header64,
            ) => image_optional_header64.section_alignment,
        };

        // 3. 判断是否可以添加节表
        let section_table_size = 40 * number_of_sections as usize;

        if size_of_headers - section_table_offset as u32 - (section_table_size as u32) < 40 {
            // 需要看一下是否可以拓展文件大小

            // 2. 计算size_of_headers对齐之后的大小 例如:0x400 -> 0x1000
            let aligned_size_of_headers =
                (size_of_headers + section_alignment - 1) & !(section_alignment - 1);
            // 可以拓展的大小: 对齐之后的大小 - size_of_headers
            let expand_size = aligned_size_of_headers - size_of_headers;
            if expand_size < 40 {
                // 如果可以拓展的大小小于40，则无法添加节表
                return Err(AppError::CannotModifyFile("无法拓展文件大小!".to_string()));
            }
            // 3. 拓展文件大小
            // 1. 打开文件
            let mut file = std::fs::OpenOptions::new()
                .read(true)
                .write(true)
                .open(self.file_path.as_ref().ok_or(AppError::NoFileOpened)?)?;
            // 设置游标到文件末尾
            file.seek(std::io::SeekFrom::End(0))?;
            // 2. 拓展文件大小
            file.set_len(file.metadata()?.len() + expand_size as u64)?;

            // 2. 映射文件
            let mut rw_mmap = unsafe {
                MmapOptions::new()
                    .map_mut(&file)
                    .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e.to_string()))?
            };
            // 循环拷贝数据 把size_of_headers之后的数据往后移动expand_size
            for i in (size_of_headers as usize..origin_file_len).rev() {
                rw_mmap[i + expand_size as usize] = rw_mmap[i];
            }

            // 把[size_of_headers, size_of_headers + expand_size]的数据清零
            for i in size_of_headers as usize..(size_of_headers as usize + expand_size as usize - 1)
            {
                rw_mmap[i] = 0;
            }
            // 修整所有节的pointer_to_raw_data
            for i in 0..number_of_sections {
                let section_header_offset = section_table_offset + 40 * i as usize;
                let pointer_to_raw_data = u32::from_le_bytes(
                    rw_mmap[section_header_offset + 20..section_header_offset + 24]
                        .try_into()
                        .unwrap(),
                );
                rw_mmap[section_header_offset + 20..section_header_offset + 24]
                    .copy_from_slice(&(pointer_to_raw_data + expand_size as u32).to_le_bytes());
            }

            // TODO: 修改SizeOfImage
            let origin_size_of_image = unsafe {
                // 在0x38的位置
                u32::from_le_bytes(
                    rw_mmap[optional_header_offset + 0x38..optional_header_offset + 0x3C]
                        .try_into()
                        .unwrap(),
                )
            };

            let add_image_size = (expand_size + section_alignment - 1) & !(section_alignment - 1);

            rw_mmap[optional_header_offset + 0x38..optional_header_offset + 0x3C]
                .copy_from_slice(&(origin_size_of_image + add_image_size as u32).to_le_bytes());

            rw_mmap.flush()?;
            // 3. 重新映射文件
            self.mmap = Some(file::mmap_file(self.file_path.as_ref().unwrap())?);
        }
        // 4. 添加节表
        // 新的节区文件偏移是文件末尾
        let new_section_offset = self.get_mmap_ref()?.len();
        let new_section_header_offset = section_table_offset + section_table_size;
        {
            // 拓展文件大小
            let mut file = std::fs::OpenOptions::new()
                .read(true)
                .write(true)
                .open(self.file_path.as_ref().ok_or(AppError::NoFileOpened)?)?;
            // 设置游标到文件末尾
            file.seek(std::io::SeekFrom::End(0))?;
            // 2. 拓展文件大小
            file.set_len(file.metadata()?.len() + section_size as u64)?;
        }
        // 写入节表数据
        let mut rw_mmap = file::mmap_file_rw(self.file_path.as_ref().unwrap())?;
        // 1. 写入节表数据
        // 1.1 写入节名
        rw_mmap[new_section_header_offset..new_section_header_offset + 8]
            .copy_from_slice(&section_name_bytes);

        // virtual_size需要对齐
        let mut virtual_size = ((section_size as u32) + section_alignment - 1) & !(section_alignment - 1);
        if virtual_size == 0 {
            virtual_size = section_alignment;
        }

        rw_mmap[new_section_header_offset + 8..new_section_header_offset + 12]
            .copy_from_slice(&(virtual_size as u32).to_le_bytes());


        // 写入节的VA
        let last_section_header_offset =
            section_table_offset + 40 * (number_of_sections - 1) as usize;
        let last_section_va = u32::from_le_bytes(
            rw_mmap[last_section_header_offset + 12..last_section_header_offset + 16]
                .try_into()
                .unwrap(),
        );
        let last_sction_size = u32::from_le_bytes(
            rw_mmap[last_section_header_offset + 16..last_section_header_offset + 20]
                .try_into()
                .unwrap(),
        );
        let last_section_alignment_size = (last_sction_size + section_alignment - 1)
            & !(section_alignment - 1);

        // 对齐后的VA, 对齐方式(last_section_va ) 和section_alignment进行对齐
        let aligned_va = ((last_section_va + last_section_alignment_size) + section_alignment - 1)
            & !(section_alignment - 1);

        rw_mmap[new_section_header_offset + 12..new_section_header_offset + 16]
            .copy_from_slice(&aligned_va.to_le_bytes());

        // 写入新节的size_of_raw_data
        rw_mmap[new_section_header_offset + 16..new_section_header_offset + 20]
            .copy_from_slice(&(section_size as u32).to_le_bytes());

        // 写入新节的PointerToRawData
        rw_mmap[new_section_header_offset + 20..new_section_header_offset + 24]
            .copy_from_slice(&(new_section_offset as u32).to_le_bytes());

        // 1.4 写入节的特征
        rw_mmap[new_section_header_offset + 36..new_section_header_offset + 40]
            .copy_from_slice(&section_characteristics.to_le_bytes());

        // 修改文件头的节表数量
        rw_mmap[file_header_offset + 2..file_header_offset + 4]
            .copy_from_slice(&(number_of_sections + 1).to_le_bytes());

        // 修改size_of_headers的大小
        let new_size_of_headers = size_of_headers + 40;
        rw_mmap[optional_header_offset + 0x3C..optional_header_offset + 0x40]
            .copy_from_slice(&new_size_of_headers.to_le_bytes());

        let origin_size_of_image = unsafe {
            // 在0x38的位置
            u32::from_le_bytes(
                rw_mmap[optional_header_offset + 0x38..optional_header_offset + 0x3C]
                    .try_into()
                    .unwrap(),
            )
        };
        let mut add_image_size =
            (section_size as u32 + section_alignment - 1) & !(section_alignment - 1);

        if add_image_size == 0 {
            add_image_size = section_alignment;
        }

        rw_mmap[optional_header_offset + 0x38..optional_header_offset + 0x3C]
            .copy_from_slice(&(origin_size_of_image + add_image_size).to_le_bytes());

        rw_mmap.flush()?;

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