use super::{
    error::{MutablePEError, PEParseError},
    header::*,
};
use std::{
    mem::{self},
    ops::{Deref, DerefMut},
};

pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
    /// Get a reference to a type at a given offset.
    fn get_ref<T>(&self, offset: usize) -> Result<&T, PEParseError> {
        let len = std::mem::size_of::<T>();
        if offset + len > self.len() {
            return Err(PEParseError::OutOfBounds);
        }
        let ptr = self.as_ptr().wrapping_offset(offset as isize) as *const T;
        Ok(unsafe { &*ptr })
    }

    // 判断是64位还是32位
    fn is_64_bit(&self) -> Result<bool, PEParseError> {
        // 先以32位加载可选头，通过可选头的Magic字段判断
        let optional_header: &ImageOptionalHeader32 = self.get_optional_header_32()?;
        let is_64_bit = optional_header.magic == 0x20b;
        Ok(is_64_bit)
    }

    // TODO: 需要一个RVA->FOA的转换函数
    // TODO: 需要一个FOA->RVA的转换函数
    // TODO: 获取数据目录、可变数据目录
    // TODO: 通过枚举获取某个数据目录的数据


    /// 将size与节对齐值进行对齐，返回对齐后的值
    /// size: 需要对齐的值
    fn align_size_with_section_alignment(&self, size: u32) -> Result<u32, PEParseError> {
        let nt_header = self.get_nt_header()?;
        let section_alignment = match nt_header {
            ImageNTHeader::NTHeader32(nt_header) => nt_header.optional_header.section_alignment,
            ImageNTHeader::NTHeader64(nt_header) => nt_header.optional_header.section_alignment,
        };
        let aligned_size = (size + section_alignment - 1) & !(section_alignment - 1);
        Ok(aligned_size)
    }

    /// 将size与文件对齐值进行对齐，返回对齐后的值
    /// size: 需要对齐的值
    fn align_size_with_file_alignment(&self, size: u32) -> Result<u32, PEParseError> {
        let nt_header = self.get_nt_header()?;
        let file_alignment = match nt_header {
            ImageNTHeader::NTHeader32(nt_header) => nt_header.optional_header.file_alignment,
            ImageNTHeader::NTHeader64(nt_header) => nt_header.optional_header.file_alignment,
        };
        let aligned_size = (size + file_alignment - 1) & !(file_alignment - 1);
        Ok(aligned_size)
    }

    /// Get the DOS header without verifying its contents.
    fn get_dos_header(&self) -> Result<&ImageDosHeader, PEParseError> {
        let result = self.get_ref::<ImageDosHeader>(0)?;
        Ok(result)
    }

    /// 作为32位的NT头解析
    fn get_nt_header32(&self) -> Result<&ImageNTHeader32, PEParseError> {
        let nt_offset = self.get_nt_headers_offset()?;
        let result = self.get_ref::<ImageNTHeader32>(nt_offset)?;
        Ok(result)
    }
    /// 作为64位的NT头解析
    fn get_nt_header64(&self) -> Result<&ImageNTHeader64, PEParseError> {
        let nt_offset = self.get_nt_headers_offset()?;
        let result = self.get_ref::<ImageNTHeader64>(nt_offset)?;
        Ok(result)
    }
    /// 动态的解析NT头 无论是32位还是64位
    fn get_nt_header(&self) -> Result<ImageNTHeader, PEParseError> {
        let is_64_bit = self.is_64_bit()?;
        let result = match is_64_bit {
            true => ImageNTHeader::NTHeader64(*self.get_nt_header64()?),
            false => ImageNTHeader::NTHeader32(*self.get_nt_header32()?),
        };
        Ok(result)
    }

    // 获取nt头的偏移
    fn get_nt_headers_offset(&self) -> Result<usize, PEParseError> {
        let dos_header = self.get_dos_header()?;
        let nt_offset = dos_header.e_lfanew.0 as usize;
        Ok(nt_offset)
    }

    // 获取文件头的偏移
    fn get_file_header_offset(&self) -> Result<usize, PEParseError> {
        // 1. 获取nt头偏移
        let nt_offset = self.get_nt_headers_offset()?;
        // 2. 计算文件头偏移
        let file_header_offset = nt_offset + std::mem::size_of::<u32>();
        Ok(file_header_offset)
    }

    // 获取文件头数据
    fn get_file_header(&self) -> Result<&ImageFileHeader, PEParseError> {
        let file_header_offset = self.get_file_header_offset()?;
        let result = self.get_ref::<ImageFileHeader>(file_header_offset)?;
        Ok(result)
    }

    // 获取节区数量
    fn get_number_of_sections(&self) -> Result<usize, PEParseError> {
        let file_header = self.get_file_header()?;
        Ok(file_header.number_of_sections as usize)
    }

    // 获取可选头大小
    fn get_size_of_optional_header(&self) -> Result<usize, PEParseError> {
        let file_header = self.get_file_header()?;
        Ok(file_header.size_of_optional_header as usize)
    }
    /// 作为32位的可选头解析
    fn get_optional_header_32(&self) -> Result<&ImageOptionalHeader32, PEParseError> {
        let file_header_offset = self.get_file_header_offset()?;
        let optional_header_offset = file_header_offset + std::mem::size_of::<ImageFileHeader>();
        let result = self.get_ref::<ImageOptionalHeader32>(optional_header_offset)?;
        Ok(result)
    }

    // 作为64位的可选头解析
    fn get_optional_header_64(&self) -> Result<&ImageOptionalHeader64, PEParseError> {
        let file_header_offset = self.get_file_header_offset()?;
        let optional_header_offset = file_header_offset + std::mem::size_of::<ImageFileHeader>();
        let result = self.get_ref::<ImageOptionalHeader64>(optional_header_offset)?;
        Ok(result)
    }
    /// 动态的解析可选头 无论是32位还是64位
    fn get_optional_header(&self) -> Result<ImageOptionalHeader, PEParseError> {
        let is_64_bit = self.is_64_bit()?;
        let result = match is_64_bit {
            true => ImageOptionalHeader::OptionalHeader64(*self.get_optional_header_64()?),
            false => ImageOptionalHeader::OptionalHeader32(*self.get_optional_header_32()?),
        };
        Ok(result)
    }
    /// 获取可选头的偏移
    fn get_optional_header_offset(&self) -> Result<usize, PEParseError> {
        let file_header_offset = self.get_file_header_offset()?;
        let optional_header_offset = file_header_offset + std::mem::size_of::<ImageFileHeader>();
        Ok(optional_header_offset)
    }

    /// 获取节区头的偏移
    fn get_section_headers_offset(&self) -> Result<usize, PEParseError> {
        // 节区头偏移在可选头之后，可选头大小是可变的，所以需要计算
        let optional_header_size = self.get_size_of_optional_header()?;
        let file_header_offset = self.get_file_header_offset()?;
        let section_header_offset =
            file_header_offset + std::mem::size_of::<ImageFileHeader>() + optional_header_size;
        Ok(section_header_offset)
    }

    // 获取所有节区头数据
    fn get_section_headers(&self) -> Result<&[ImageSectionHeader], PEParseError> {
        // 1. 获取节区数量
        let number_of_sections = self.get_number_of_sections()?;
        // 2. 获取节区头偏移
        let section_header_offset = self.get_section_headers_offset()?;
        // 3. 获取节区头数据
        unsafe {
            let ptr = self
                .as_ptr()
                .wrapping_offset(section_header_offset as isize)
                as *const ImageSectionHeader;
            let result = std::slice::from_raw_parts(ptr, number_of_sections);
            Ok(result)
        }
    }

    /// 获取头映射最大可扩容大小
    fn get_header_max_expand_size(&self) -> Result<u32, PEParseError> {
        // 最大可扩容大小 = 节对齐大小 - 最后一个节区头最后一个字节的偏移
        let nt_header: ImageNTHeader = self.get_nt_header()?;
        let (section_alignment, file_alignment) = match nt_header {
            ImageNTHeader::NTHeader32(nt_header) => (
                nt_header.optional_header.section_alignment,
                nt_header.optional_header.file_alignment,
            ),
            ImageNTHeader::NTHeader64(nt_header) => (
                nt_header.optional_header.section_alignment,
                nt_header.optional_header.file_alignment,
            ),
        };
        let section_header_offset = self.get_section_headers_offset()?;
        let section_headers = self.get_number_of_sections()? as u32;
        let last_section_header_offset = section_header_offset
            + (section_headers as usize * std::mem::size_of::<ImageSectionHeader>());
        let mut max_size_of_headers = section_alignment - last_section_header_offset as u32;
        // max_size_of_headers 向下对齐
        max_size_of_headers &= !(file_alignment - 1);

        Ok(max_size_of_headers)
    }

    /// 获取头部空余空间大小
    /// 头部空余空间大小 = 可选头的SizeOfHeaders - 最后一个节区头最后一个字节的偏移
    fn get_header_empty_space_size(&self) -> Result<u32, PEParseError> {
        let nt_header: ImageNTHeader = self.get_nt_header()?;
        let size_of_headers = match nt_header {
            ImageNTHeader::NTHeader32(nt_header) => nt_header.optional_header.size_of_headers,
            ImageNTHeader::NTHeader64(nt_header) => nt_header.optional_header.size_of_headers,
        };
        let section_header_offset = self.get_section_headers_offset()?;
        let section_headers = self.get_number_of_sections()? as u32;
        let last_section_header_offset = section_header_offset
            + (section_headers as usize * std::mem::size_of::<ImageSectionHeader>());
        let empty_space_size = size_of_headers - last_section_header_offset as u32;
        Ok(empty_space_size)
    }
}
/// 可修改的PE trait
pub trait MutablePE: ReadOnlyPE + DerefMut<Target = [u8]> + AsMut<[u8]> {
    /// 实现一个get_mut
    fn get_mut<T>(&mut self, offset: usize) -> Result<&mut T, PEParseError> {
        let len = std::mem::size_of::<T>();
        if offset + len > self.len() {
            return Err(PEParseError::OutOfBounds);
        }
        let ptr = self.as_mut_ptr().wrapping_offset(offset as isize) as *mut T;
        Ok(unsafe { &mut *ptr })
    }

    /// 获取可变的文件头数据
    fn get_file_header_mut(&mut self) -> Result<&mut ImageFileHeader, PEParseError> {
        let file_header_offset = self.get_file_header_offset()?;
        let result = self.get_mut::<ImageFileHeader>(file_header_offset)?;
        Ok(result)
    }
    /// 获取可变的节区头数据
    fn get_section_headers_mut(&mut self) -> Result<&mut [ImageSectionHeader], PEParseError> {
        // 1. 获取节区数量
        let number_of_sections = self.get_number_of_sections()?;
        // 2. 获取节区头偏移
        let section_header_offset = self.get_section_headers_offset()?;
        // 3. 获取节区头数据
        unsafe {
            let ptr = self
                .as_mut_ptr()
                .wrapping_offset(section_header_offset as isize)
                as *mut ImageSectionHeader;
            let result = std::slice::from_raw_parts_mut(ptr, number_of_sections);
            Ok(result)
        }
    }
    /// 获取可变的NT头数据64位
    fn get_nt_header64_mut(&mut self) -> Result<&mut ImageNTHeader64, PEParseError> {
        let nt_offset = self.get_nt_headers_offset()?;
        let result = self.get_mut::<ImageNTHeader64>(nt_offset)?;
        Ok(result)
    }

    /// 获取可变的NT头数据32位
    fn get_nt_header32_mut(&mut self) -> Result<&mut ImageNTHeader32, PEParseError> {
        let nt_offset = self.get_nt_headers_offset()?;
        let result = self.get_mut::<ImageNTHeader32>(nt_offset)?;
        Ok(result)
    }

    /// 获取可变的NT头数据
    fn get_nt_header_mut(&mut self) -> Result<ImageNTHeaderMut, PEParseError> {
        let is_64_bit = self.is_64_bit()?;
        let result = match is_64_bit {
            true => ImageNTHeaderMut::NTHeader64(self.get_nt_header64_mut()?),
            false => ImageNTHeaderMut::NTHeader32(self.get_nt_header32_mut()?),
        };
        Ok(result)
    }

    /// 扩大头映射大小，调用者必须要保证当挪动所有节区数据时，不会超出文件范围
    /// 也就是，调用者必须要在扩大头映射大小之前，先扩大文件大小
    /// add_size: 需要扩大的大小，不需要进行对齐，会自动对齐
    fn expand_headers(&mut self, add_size: u32) -> Result<(), MutablePEError> {
        // aligned_size: 需要增加的大小，对齐后的值
        let aligned_file_size = self.align_size_with_file_alignment(add_size)?;
        let max_expand_size = self.get_header_max_expand_size()?;
        let optional_header = self.get_optional_header()?;
        let origin_size_for_headers = optional_header.get_size_of_headers();
        // 如果扩大后超过了最大可扩容大小，就不允许扩容
        if aligned_file_size > max_expand_size {
            return Err(MutablePEError::CannotExpandHeader(
                aligned_file_size,
                origin_size_for_headers,
            ));
        }
        // 检查一下把，万一调用者没有扩大文件大小
        let sections = self.get_section_headers()?;
        let self_len = self.len();
        // 需要反向遍历，找到最后一个size_of_raw_data不为0的节区
        // 然后计算这个节区的PointerToRawData+（SizeOfRawData对齐到FileAlignment的值）
        let mut last_section_end = 0;
        for section in sections.iter().rev() {
            if section.size_of_raw_data != 0 {
                last_section_end = section.pointer_to_raw_data
                    + self.align_size_with_file_alignment(section.size_of_raw_data)?;
                if last_section_end as usize + aligned_file_size as usize > self_len {
                    return Err(MutablePEError::CannotExpandFileSize(
                        last_section_end as u32 + aligned_file_size,
                        self_len as u32,
                    ));
                }
                break;
            }
        }
        let mut nt_header = self.get_nt_header_mut()?;
        let op = nt_header.get_optional_header_mut();
        // 更新可选头的SizeOfHeaders和SizeOfImage
        match op {
            ImageOptionalHeaderMut::OptionalHeader32(op) => {
                op.size_of_headers += aligned_file_size;
            }
            ImageOptionalHeaderMut::OptionalHeader64(op) => {
                op.size_of_headers += aligned_file_size;
            }
        }
        let mut_section = self.get_section_headers_mut()?;
        let start_offset = mut_section.first().unwrap().pointer_to_raw_data;
        let end_offset = last_section_end;
        // 更新节区头的PointerToRawData
        for section in mut_section.iter_mut() {
            section.pointer_to_raw_data += aligned_file_size;
        }
        {
            let self_data = self.as_mut();
            // 把[start_offset, end_offset)的数据往后挪动aligned_file_size
            self_data.copy_within(
                start_offset as usize..end_offset as usize,
                (start_offset + aligned_file_size) as usize,
            );
            // 把[start_offset, start_offset+aligned_file_size)的数据填充为0
            self_data[start_offset as usize..(start_offset + aligned_file_size) as usize].fill(0);
        }
        Ok(())
    }

    /// 添加一个节区
    /// section_name: 节区名
    /// section_size: 节区大小
    /// section_characteristics: 节区特征
    fn add_section(
        &mut self,
        section_name: &[u8; 8],
        section_size: u32,
        section_characteristics: SectionCharacteristics,
    ) -> Result<(), MutablePEError> {
        // 1. 判断是否有空余空间
        let empty_space_size = self.get_header_empty_space_size()?;
        const SIZE_OF_SECTION_HEADER: u32 = mem::size_of::<ImageSectionHeader>() as u32;
        // section_size与文件对齐后的大小
        let file_alignment = self.align_size_with_file_alignment(section_size)?;

        if empty_space_size < SIZE_OF_SECTION_HEADER {
            // 判断是否能够扩容
            let max_expand_size = self.get_header_max_expand_size()?;
            if max_expand_size < file_alignment {
                return Err(MutablePEError::CannotExpandHeader(
                    file_alignment,
                    max_expand_size,
                ));
            }
            self.expand_headers(file_alignment)?;
        }
        // 增加节区头
        let sections_offset = self.get_section_headers_offset()?;
        let number_of_sections = self.get_number_of_sections()?;
        let new_section_offset =
            sections_offset + (number_of_sections * SIZE_OF_SECTION_HEADER as usize) as usize;
        let new_section = unsafe {
            let ptr = self
                .as_mut_ptr()
                .wrapping_offset(new_section_offset as isize)
                as *mut ImageSectionHeader;
            &mut *ptr
        };
        new_section.name = *section_name;
        // 新节区的virtual_size是section_size对齐到section_alignment的值
        let mut section_alignment = self.align_size_with_section_alignment(section_size)?;
        if section_size == 0 {
            section_alignment = self.align_size_with_section_alignment(1)?;
        }
        new_section.virtual_size = section_alignment;
        // 如果section_size是0，那么virtual_size必须得是最小的section_alignment
        // 新节区的virtual_address是上一个节区的virtual_address+(pointer_to_raw_data 对齐到 section_alignment的值)
        let last_section = self.get_section_headers()?.last().unwrap();
        let mut last_section_size_of_raw_data_aligned =
            self.align_size_with_section_alignment(last_section.size_of_raw_data)?;
        if last_section_size_of_raw_data_aligned == 0 {
            last_section_size_of_raw_data_aligned = self.align_size_with_section_alignment(1)?;
        }
        new_section.virtual_address =
            last_section.virtual_address + last_section_size_of_raw_data_aligned;
        new_section.size_of_raw_data = file_alignment;
        // 如果上一个节区的size_of_raw_data是0，那么上一个节区的pointer_to_raw_data是可以乱写的
        // 正确的做法应该是，找到上一个size_of_raw_data不为0的节区，然后计算
        let sections = self.get_section_headers()?;
        // 反向遍历
        for section in sections.iter().rev() {
            if section.size_of_raw_data != 0 {
                new_section.pointer_to_raw_data = section.pointer_to_raw_data
                    + self.align_size_with_file_alignment(section.size_of_raw_data)?;
                break;
            }
        }
        new_section.characteristics = section_characteristics;
        // 其他字段最好填充为0
        new_section.number_of_linenumbers = 0;
        new_section.number_of_relocations = 0;
        new_section.pointer_to_linenumbers = 0;
        new_section.pointer_to_relocations = 0;
        // 更新文件头的NumberOfSections
        let file_header = self.get_file_header_mut()?;
        file_header.number_of_sections += 1;

        // 更新可选头的SizeOfImage
        let mut nt_header = self.get_nt_header_mut()?;
        let op = nt_header.get_optional_header_mut();
        match op {
            ImageOptionalHeaderMut::OptionalHeader32(op) => {
                op.size_of_image += section_alignment;
            }
            ImageOptionalHeaderMut::OptionalHeader64(op) => {
                op.size_of_image += section_alignment;
            }
        }
        Ok(())
    }
}
