use super::types::*;
use bitflags::bitflags;
use serde::Serialize;

#[repr(C)]
#[derive(Debug, Clone, Copy, Serialize)]
pub struct ImageDosHeader {
    pub e_magic: u16, // Magic number 固定值 0x5A4D
    pub e_cblp: u16,
    pub e_cp: u16,
    pub e_crlc: u16,
    pub e_cparhdr: u16,
    pub e_minalloc: u16,
    pub e_maxalloc: u16,
    pub e_ss: u16,
    pub e_sp: u16,
    pub e_csum: u16,
    pub e_ip: u16,
    pub e_cs: u16,
    pub e_lfarlc: u16,
    pub e_ovno: u16,
    pub e_res: [u16; 4],
    pub e_oemid: u16,
    pub e_oeminfo: u16,
    pub e_res2: [u16; 10],
    pub e_lfanew: Offset, // File address of new exe header nt头的偏移
}

#[repr(C)]
#[derive(Debug, Clone, Copy, Serialize)]
pub struct ImageNTHeader32 {
    pub signature: u32,
    pub file_header: ImageFileHeader,
    pub optional_header: ImageOptionalHeader32,
}

#[repr(C)]
#[derive(Debug, Clone, Copy, Serialize)]
pub struct ImageNTHeader64 {
    pub signature: u32,
    pub file_header: ImageFileHeader,
    pub optional_header: ImageOptionalHeader64,
}

#[repr(C)]
#[derive(Serialize, Clone, Copy)]
#[serde(untagged)]
pub enum ImageNTHeader {
    NTHeader32(ImageNTHeader32),
    NTHeader64(ImageNTHeader64),
}

#[repr(C)]
#[derive(Serialize)]
#[serde(untagged)]
pub enum ImageNTHeaderMut<'a> {
    NTHeader32(&'a mut ImageNTHeader32),
    NTHeader64(&'a mut ImageNTHeader64),
}

impl<'a> ImageNTHeaderMut<'a> {
    pub fn get_optional_header_mut(&mut self) -> ImageOptionalHeaderMut {
        match self {
            ImageNTHeaderMut::NTHeader32(header) => {
                ImageOptionalHeaderMut::OptionalHeader32(&mut header.optional_header)
            }
            ImageNTHeaderMut::NTHeader64(header) => {
                ImageOptionalHeaderMut::OptionalHeader64(&mut header.optional_header)
            }
        }
    }
}

#[repr(C)]
#[derive(Debug, Clone, Copy, Serialize)]
pub struct ImageFileHeader {
    pub machine: u16,
    pub number_of_sections: u16,
    pub time_date_stamp: u32,
    pub pointer_to_symbol_table: Offset,
    pub number_of_symbols: u32,
    pub size_of_optional_header: u16,
    pub characteristics: FileCharacteristics,
}
impl Serialize for FileCharacteristics {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        // 直接返回bitflags的整数值
        serializer.serialize_u16(self.bits())
    }
}

impl Serialize for DLLCharacteristics {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.serialize_u16(self.bits())
    }
}

impl Serialize for SectionCharacteristics {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.serialize_u32(self.bits())
    }
}

bitflags! {
    #[repr(C)]
    #[derive(Debug, Clone, Copy)]
    pub struct FileCharacteristics: u16 {
        const RELOCS_STRIPPED         = 0x0001;
        const EXECUTABLE_IMAGE        = 0x0002;
        const LINE_NUMS_STRIPPED      = 0x0004;
        const LOCAL_SYMS_STRIPPED     = 0x0008;
        const AGGRESSIVE_WS_TRIM      = 0x0010;
        const LARGE_ADDRESS_AWARE     = 0x0020;
        const BYTES_REVERSED_LO       = 0x0080;
        const MACHINE_32BIT           = 0x0100;
        const DEBUG_STRIPPED          = 0x0200;
        const REMOVABLE_RUN_FROM_SWAP = 0x0400;
        const NET_RUN_FROM_SWAP       = 0x0800;
        const SYSTEM                  = 0x1000;
        const DLL                     = 0x2000;
        const UP_SYSTEM_ONLY          = 0x4000;
        const BYTES_REVERSED_HI       = 0x8000;
    }
    #[repr(C)]
    #[derive(Debug, Clone, Copy)]
    pub struct DLLCharacteristics: u16 {
        const RESERVED1             = 0x0001;
        const RESERVED2             = 0x0002;
        const RESERVED4             = 0x0004;
        const RESERVED8             = 0x0008;
        const HIGH_ENTROPY_VA       = 0x0020;
        const DYNAMIC_BASE          = 0x0040;
        const FORCE_INTEGRITY       = 0x0080;
        const NX_COMPAT             = 0x0100;
        const NO_ISOLATION          = 0x0200;
        const NO_SEH                = 0x0400;
        const NO_BIND               = 0x0800;
        const APPCONTAINER          = 0x1000;
        const WDM_DRIVER            = 0x2000;
        const GUARD_CF              = 0x4000;
        const TERMINAL_SERVER_AWARE = 0x8000;
    }
    #[repr(C)]
    #[derive(Debug, Clone, Copy)]
    pub struct SectionCharacteristics: u32 {
        /// Reserved for future use.
        const TYPE_REG               = 0x00000000;
        /// Reserved for future use.
        const TYPE_DSECT             = 0x00000001;
        /// Reserved for future use.
        const TYPE_NOLOAD            = 0x00000002;
        /// Reserved for future use.
        const TYPE_GROUP             = 0x00000004;
        /// The section should not be padded to the next boundary.
        /// This flag is obsolete and is replaced by IMAGE_SCN_ALIGN_1BYTES.
        /// This is valid only for object files.
        const TYPE_NO_PAD            = 0x00000008;
        /// Reserved for future use.
        const TYPE_COPY              = 0x00000010;
        /// The section contains executable code.
        const CNT_CODE               = 0x00000020;
        /// The section contains initialized data.
        const CNT_INITIALIZED_DATA   = 0x00000040;
        /// The section contains uninitialized data.
        const CNT_UNINITIALIZED_DATA = 0x00000080;
        /// Reserved for future use.
        const LNK_OTHER              = 0x00000100;
        /// The section contains comments or other information.
        /// The .drectve section has this type. This is valid for object files only.
        const LNK_INFO               = 0x00000200;
        /// Reserved for future use.
        const TYPE_OVER              = 0x00000400;
        /// The section will not become part of the image.
        /// This is valid only for object files.
        const LNK_REMOVE             = 0x00000800;
        /// The section contains COMDAT data. This is valid only for object files.
        const LNK_COMDAT             = 0x00001000;
        /// Unknown/Reserved.
        const RESERVED               = 0x00002000;
        /// Unknown flag.
        const MEM_PROTECTED          = 0x00004000;
        /// Unknown flag.
        const NO_DEFER_SPEC_EXC      = 0x00004000;
        /// The section contains data referenced through the global pointer (GP).
        const GPREL                  = 0x00008000;
        /// Reserved for future use.
        const MEM_FARDATA            = 0x00008000;
        /// Reserved for future use.
        const MEM_SYSHEAP            = 0x00010000;
        /// Reserved for future use.
        const MEM_PURGEABLE          = 0x00020000;
        /// Reserved for future use.
        const MEM_16BIT              = 0x00020000;
        /// Reserved for future use.
        const MEM_LOCKED             = 0x00040000;
        /// Reserved for future use.
        const MEM_PRELOAD            = 0x00080000;
        /// Align data on a 1-byte boundary. Valid only for object files.
        const ALIGN_1BYTES           = 0x00100000;
        /// Align data on a 2-byte boundary. Valid only for object files.
        const ALIGN_2BYTES           = 0x00200000;
        /// Align data on a 4-byte boundary. Valid only for object files.
        const ALIGN_4BYTES           = 0x00300000;
        /// Align data on an 8-byte boundary. Valid only for object files.
        const ALIGN_8BYTES           = 0x00400000;
        /// Align data on a 16-byte boundary. Valid only for object files.
        const ALIGN_16BYTES          = 0x00500000;
        /// Align data on a 32-byte boundary. Valid only for object files.
        const ALIGN_32BYTES          = 0x00600000;
        /// Align data on a 64-byte boundary. Valid only for object files.
        const ALIGN_64BYTES          = 0x00700000;
        /// Align data on a 128-byte boundary. Valid only for object files.
        const ALIGN_128BYTES         = 0x00800000;
        /// Align data on a 256-byte boundary. Valid only for object files.
        const ALIGN_256BYTES         = 0x00900000;
        /// Align data on a 512-byte boundary. Valid only for object files.
        const ALIGN_512BYTES         = 0x00A00000;
        /// Align data on a 1024-byte boundary. Valid only for object files.
        const ALIGN_1024BYTES        = 0x00B00000;
        /// Align data on a 2048-byte boundary. Valid only for object files.
        const ALIGN_2048BYTES        = 0x00C00000;
        /// Align data on a 4096-byte boundary. Valid only for object files.
        const ALIGN_4096BYTES        = 0x00D00000;
        /// Align data on an 8192-byte boundary. Valid only for object files.
        const ALIGN_8192BYTES        = 0x00E00000;
        /// Mask for alignment.
        const ALIGN_MASK             = 0x00F00000;
        /// The section contains extended relocations.
        const LNK_NRELOC_OVFL        = 0x01000000;
        /// The section can be discarded as needed.
        const MEM_DISCARDABLE        = 0x02000000;
        /// The section cannot be cached.
        const MEM_NOT_CACHED         = 0x04000000;
        /// The section is not pageable.
        const MEM_NOT_PAGED          = 0x08000000;
        /// The section can be shared in memory.
        const MEM_SHARED             = 0x10000000;
        /// The section can be executed as code.
        const MEM_EXECUTE            = 0x20000000;
        /// The section can be read.
        const MEM_READ               = 0x40000000;
        /// The section can be written to.
        const MEM_WRITE              = 0x80000000;
    }
}
#[repr(C)]
#[derive(Debug, Clone, Copy, Serialize)]
pub struct ImageOptionalHeader32 {
    pub magic: u16,
    pub major_linker_version: u8,
    pub minor_linker_version: u8,
    pub size_of_code: u32,
    pub size_of_initialized_data: u32,
    pub size_of_uninitialized_data: u32,
    pub address_of_entry_point: u32,
    pub base_of_code: u32,
    pub base_of_data: u32,
    pub image_base: u32,
    pub section_alignment: u32,
    pub file_alignment: u32,
    pub major_operating_system_version: u16,
    pub minor_operating_system_version: u16,
    pub major_image_version: u16,
    pub minor_image_version: u16,
    pub major_subsystem_version: u16,
    pub minor_subsystem_version: u16,
    pub win32_version_value: u32,
    pub size_of_image: u32,
    pub size_of_headers: u32,
    pub checksum: u32,
    pub subsystem: u16,
    pub dll_characteristics: DLLCharacteristics,
    pub size_of_stack_reserve: u32,
    pub size_of_stack_commit: u32,
    pub size_of_heap_reserve: u32,
    pub size_of_heap_commit: u32,
    pub loader_flags: u32,
    pub number_of_rva_and_sizes: u32,
}

#[repr(C)]
#[derive(Copy, Clone, Debug, Serialize)]
pub struct ImageOptionalHeader64 {
    pub magic: u16,
    pub major_linker_version: u8,
    pub minor_linker_version: u8,
    pub size_of_code: u32,
    pub size_of_initialized_data: u32,
    pub size_of_uninitialized_data: u32,
    pub address_of_entry_point: u32,
    pub base_of_code: u32,
    pub image_base: u64,
    pub section_alignment: u32,
    pub file_alignment: u32,
    pub major_operating_system_version: u16,
    pub minor_operating_system_version: u16,
    pub major_image_version: u16,
    pub minor_image_version: u16,
    pub major_subsystem_version: u16,
    pub minor_subsystem_version: u16,
    pub win32_version_value: u32,
    pub size_of_image: u32,
    pub size_of_headers: u32,
    pub checksum: u32,
    pub subsystem: u16,
    pub dll_characteristics: DLLCharacteristics,
    pub size_of_stack_reserve: u64,
    pub size_of_stack_commit: u64,
    pub size_of_heap_reserve: u64,
    pub size_of_heap_commit: u64,
    pub loader_flags: u32,
    pub number_of_rva_and_sizes: u32,
}
#[repr(C)]
#[derive(Serialize, Clone, Copy)]
#[serde(untagged)]
pub enum ImageOptionalHeader {
    OptionalHeader32(ImageOptionalHeader32),
    OptionalHeader64(ImageOptionalHeader64),
}

#[repr(C)]
#[derive(Copy, Clone, Eq, PartialEq, Default, Debug)]
pub struct ImageDataDirectory {
    pub virtual_address: u32,
    pub size: u32,
}

#[derive(Copy, Clone, Eq, PartialEq, Debug)]
pub enum ImageDirectoryEntry {
    Export = 0,
    Import = 1,
    Resource = 2,
    Exception = 3,
    Security = 4,
    BaseReloc = 5,
    Debug = 6,
    Architecture = 7,
    GlobalPTR = 8,
    TLS = 9,
    LoadConfig = 10,
    BoundImport = 11,
    IAT = 12,
    DelayImport = 13,
    COMDescriptor = 14,
    Reserved = 15,
}

#[repr(C)]
#[derive(Serialize)]
#[serde(untagged)]
pub enum ImageOptionalHeaderMut<'a> {
    OptionalHeader32(&'a mut ImageOptionalHeader32),
    OptionalHeader64(&'a mut ImageOptionalHeader64),
}

impl ImageOptionalHeader {
    pub fn get_size_of_headers(&self) -> u32 {
        match self {
            ImageOptionalHeader::OptionalHeader32(header) => header.size_of_headers,
            ImageOptionalHeader::OptionalHeader64(header) => header.size_of_headers,
        }
    }

    pub fn get_size_of_image(&self) -> u32 {
        match self {
            ImageOptionalHeader::OptionalHeader32(header) => header.size_of_image,
            ImageOptionalHeader::OptionalHeader64(header) => header.size_of_image,
        }
    }
    pub fn get_size_of_optional_header(&self) -> u32 {
        match self {
            ImageOptionalHeader::OptionalHeader32(_) => {
                size_of::<ImageOptionalHeader32>() as u32
            }
            ImageOptionalHeader::OptionalHeader64(_) => {
                size_of::<ImageOptionalHeader64>() as u32
            }
        }
    }
    /// 获取数据目录数量
    pub fn get_data_directory_count(&self) -> u32 {
        // 如果数据目录数量大于16，则返回16
        match self {
            ImageOptionalHeader::OptionalHeader32(header) => header.number_of_rva_and_sizes,
            ImageOptionalHeader::OptionalHeader64(header) => header.number_of_rva_and_sizes,
        }
        .min(16)
    }
}

#[repr(C)]
#[derive(Debug, Clone, Copy, Serialize)]
pub struct ImageSectionHeader {
    pub name: [u8; 8],
    pub virtual_size: u32,
    pub virtual_address: u32,
    pub size_of_raw_data: u32,
    pub pointer_to_raw_data: u32,
    pub pointer_to_relocations: u32,
    pub pointer_to_linenumbers: u32,
    pub number_of_relocations: u16,
    pub number_of_linenumbers: u16,
    pub characteristics: SectionCharacteristics,
}

#[repr(C)]
#[derive(Copy, Clone, Eq, PartialEq, Debug, Default)]
pub struct ImageImportDescriptor {
    pub original_first_thunk: u32,
    pub time_date_stamp: u32,
    pub forwarder_chain: u32,
    pub name: u32,
    pub first_thunk: u32,
}

bitflags! {
    /// 导入查找表 32位
    /// [Import Lookup Table](https://learn.microsoft.com/en-us/windows/win32/debug/pe-format#import-lookup-table)
    #[derive(Debug, Clone, Copy, Serialize)]
    pub struct ImportLookupTable32: u32 {
        /// 第31位为1表示导入函数通过序号导入
        const IMPORT_BY_ORDINAL = 0x80000000;
        /// 15-0 位表示导入函数的序号
        const IMPORT_BY_ORDINAL_MASK = 0x0000FFFF;
        /// 30-0 位表示导入函数的RVA
        const IMPORT_BY_RVA_MASK = 0x7FFFFFFF;
    }

    /// 导入查找表 64位
    /// [Import Lookup Table](https://learn.microsoft.com/en-us/windows/win32/debug/pe-format#import-lookup-table)
    #[derive(Debug, Clone, Copy, Serialize)]
    pub struct ImportLookupTable64: u64 {
        /// 第63位为1表示导入函数通过序号导入
        const IMPORT_BY_ORDINAL = 0x8000000000000000;
        /// 15-0 位表示导入函数的序号
        const IMPORT_BY_ORDINAL_MASK = 0x0000FFFF;
        /// 30-0 位表示导入函数的RVA
        const IMPORT_BY_RVA_MASK = 0x7FFFFFFFFFFFFFFF;
        /// 62-31 must be zero.
        const RESERVED = 0x7FFFFFFF00000000;
    }

}