use std::{
    alloc::{alloc, Layout},
    io::{Seek, Write},
};

use crate::app_error::AppError;

use pe_parse::{header::ImageDosHeader, pe::ReadOnlyPE};

use memmap2::Mmap;
use sysinfo::{self, System};
use windows::{
    core::PWSTR,
    Win32::{
        Foundation::{CloseHandle, HMODULE},
        System::{
            Diagnostics::Debug::ReadProcessMemory,
            ProcessStatus::EnumProcessModules,
            Threading::{
                OpenProcess, QueryFullProcessImageNameW, PROCESS_ACCESS_RIGHTS, PROCESS_ALL_ACCESS,
                PROCESS_NAME_WIN32,
            },
        },
    },
};

/// 遍历进程并返回进程名\PID
pub fn get_process_list() -> Vec<(String, u32)> {
    let mut sys = System::new_all();
    sys.refresh_all();
    let mut process_list = Vec::new();
    for (pid, process) in sys.processes() {
        let pid = pid.as_u32();
        process_list.push((process.name().to_str().unwrap().to_string(), pid));
    }
    process_list
}
pub fn dump_process(pid: u32, save_file_path: &str) -> Result<(), AppError> {
    #[warn(unused_assignments)]
    let mut file_path = String::new();
    // 1. 打开进程

    let mut image_base = 0; // 模块基址
    let h_process = unsafe {
        OpenProcess(PROCESS_ACCESS_RIGHTS(PROCESS_ALL_ACCESS.0), false, pid)
            .map_err(|e| AppError::OpenProcessError(e.to_string()))?
    };
    // 2. 获取进程对应的文件路径
    let mut buffer = [0u16; 1024];
    let pwstr = PWSTR(buffer.as_mut_ptr());
    let mut size = buffer.len() as u32;
    unsafe {
        let result = QueryFullProcessImageNameW(h_process, PROCESS_NAME_WIN32, pwstr, &mut size);
        if result.is_err() {
            return Err(AppError::QueryFullProcessImageNameWError(
                result.unwrap_err().to_string(),
            ));
        }
    }
    unsafe {
        file_path = pwstr.to_string().unwrap();
        println!("进程路径: {:?}", file_path);
    }
    // 创建新文件
    let mut save_file = std::fs::OpenOptions::new()
        .write(true)
        .create(true)
        .open(save_file_path)
        .map_err(|e| AppError::Io(e))?;

    // 3. 读取进程的文件，读取它的PE头
    // 3.1 映射文件
    let file = std::fs::OpenOptions::new()
        .read(true)
        .open(file_path)
        .map_err(|e| AppError::Io(e))?;
    let mmap = unsafe { Mmap::map(&file).map_err(|e| AppError::Io(e)) }?;
    // 3.2 读取PE头
    {
        // 3.2.1 读取DOS头
        let dos_header = mmap.get_dos_header()?;
        let write_size = std::mem::size_of::<ImageDosHeader>() as u64;
        save_file.write_all(&mmap[..write_size as usize])?;
        // 3.2.2 读取PE头
        let nt_header = mmap.get_nt_header()?;
        let nt_header_size = match nt_header {
            pe_parse::header::ImageNTHeader::NTHeader32(_) => {
                std::mem::size_of::<pe_parse::header::ImageNTHeader32>() as u64
            }
            pe_parse::header::ImageNTHeader::NTHeader64(_) => {
                std::mem::size_of::<pe_parse::header::ImageNTHeader64>() as u64
            }
        };
        let nt_header_offset = dos_header.e_lfanew;
        // 移动文件指针
        save_file.seek(std::io::SeekFrom::Start(nt_header_offset.0 as u64))?;
        // 写入PE头
        save_file.write_all(
            &mmap[nt_header_offset.0 as usize
                ..(nt_header_offset.0 as usize + nt_header_size as usize) as usize],
        )?;
        // 写入nt_header_offset.0 - sizeof(dos_header)的数据
        // 移动文件指针到dos_header之后
        save_file.seek(std::io::SeekFrom::Start(write_size))?;
        save_file.write_all(&mmap[write_size as usize..nt_header_offset.0 as usize])?;
        // 移动文件指针到最后
        save_file.seek(std::io::SeekFrom::Start(
            nt_header_offset.0 as u64 + nt_header_size,
        ))?;
        // 数据目录
        let data_directory_offset = nt_header_offset.0 as usize + nt_header_size as usize;
        save_file.write_all(&mmap[data_directory_offset..data_directory_offset + 16 * 17])?;
        // 节区头
        let sections_num = mmap.get_section_headers()?.len();
        let section_header_offset = mmap.get_section_headers_offset()?;
        let section_header_size = std::mem::size_of::<pe_parse::header::ImageSectionHeader>()
            as u64
            * sections_num as u64;
        save_file.seek(std::io::SeekFrom::Start(section_header_offset as u64))?;
        save_file.write_all(
            &mmap[section_header_offset as usize
                ..(section_header_offset as usize + section_header_size as usize) as usize],
        )?;
        save_file.flush()?;
    }
    // 获取模块基址
    {
        let mut h_module = HMODULE::default();
        let mut cb_needed = 0;
        unsafe {
            EnumProcessModules(
                h_process,
                &mut h_module,
                std::mem::size_of::<HMODULE>() as u32,
                &mut cb_needed,
            )
            .unwrap();
            // 获取模块基址
            image_base = h_module.0 as u64;
        }
    }

    // 5. 遍历节表头，从进程中读取节的内容，写入文件
    {
        let sections = mmap.get_section_headers()?;
        for section in sections.iter() {
            let section_va = section.virtual_address as u64;
            let section_size = section.size_of_raw_data as u64;
            let section_size_align = mmap.align_size_with_section_alignment(section_size as u32)?;
            let section_offset = image_base + section_va;
            let section_foa = section.pointer_to_raw_data as u64;

            let data_buf = unsafe {
                let layout = Layout::array::<u8>(section_size_align as usize).unwrap();
                alloc(layout)
            };

            let mut read_size = 0;

            unsafe {
                ReadProcessMemory(
                    h_process,
                    section_offset as *const _,
                    data_buf as *mut _,
                    section_size_align as usize,
                    Some(&mut read_size),
                )
                .unwrap();
            }
            println!(
                "读取节: VA: {:#x}, Size: {:#x}, FOA: {:#x}",
                section_va, section_size, section_foa
            );
            println!(
                "读取节: Size: {:#x}, ReadSize: {:#x}",
                section_size_align, read_size
            );
            // 如果文件长度不够，则先扩容
            let file_size = save_file.metadata().unwrap().len();
            if file_size < section_foa + section_size_align as u64 {
                // 文件指针要移动到末尾
                save_file.seek(std::io::SeekFrom::End(0))?;
                save_file.set_len(section_foa + section_size_align as u64)?;
            }

            save_file.seek(std::io::SeekFrom::Start(section_foa))?;
            save_file.write_all(unsafe {
                std::slice::from_raw_parts(data_buf as *const u8, section_size_align as usize)
            })?;
        }
    }
    // 6. 关闭进程
    unsafe {
        CloseHandle(h_process).unwrap();
    };
    // 7. 返回
    Ok(())
}

// 来个测试
#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_get_process_list() {
        let process_list = get_process_list();
        for (name, pid) in process_list {
            println!("进程名: {}, PID: {}", name, pid);
        }
    }
    #[test]
    fn test_dump_process() {
        dump_process(11696, "D:\\download\\a.exe").unwrap();
    }
}
