/// IAT (Import Address Table) patching.
///
/// We define the PE structs ourselves rather than pulling in windows-sys PE
/// headers (which are in a different crate / feature path).
use windows_sys::Win32::{
    Foundation::HMODULE,
    System::{
        LibraryLoader::{GetModuleHandleA, GetProcAddress},
        Memory::{VirtualProtect, PAGE_EXECUTE_READWRITE},
    },
};

// ── Minimal PE struct definitions ─────────────────────────────────────────────

#[repr(C)]
struct ImageDosHeader {
    e_magic: u16,
    _pad: [u16; 29],
    e_lfanew: i32,
}

#[repr(C)]
struct ImageFileHeader {
    machine: u16,
    number_of_sections: u16,
    time_date_stamp: u32,
    pointer_to_symbol_table: u32,
    number_of_symbols: u32,
    size_of_optional_header: u16,
    characteristics: u16,
}

#[repr(C)]
struct ImageDataDirectory {
    virtual_address: u32,
    size: u32,
}

#[repr(C)]
struct ImageOptionalHeader64 {
    magic: u16,
    _pad: [u8; 110],
    data_directory: [ImageDataDirectory; 16],
}

#[repr(C)]
struct ImageNtHeaders64 {
    signature: u32,
    file_header: ImageFileHeader,
    optional_header: ImageOptionalHeader64,
}

#[repr(C)]
struct ImageImportDescriptor {
    original_first_thunk: u32,
    time_date_stamp: u32,
    forwarder_chain: u32,
    name: u32,
    first_thunk: u32,
}

// ── IAT patching ──────────────────────────────────────────────────────────────

/// Replace every IAT slot in the main module that currently holds
/// `original_fn` with `hook_fn`.
///
/// # Safety
/// Caller must ensure hook_fn has the same calling convention and signature.
pub unsafe fn patch_iat(original_fn: *const (), hook_fn: *const ()) -> usize {
    let module = GetModuleHandleA(std::ptr::null());
    patch_module(module, original_fn, hook_fn)
}

/// Patch the IAT of a specific already-loaded DLL (e.g. b"EAWebKit.dll\0").
pub unsafe fn patch_iat_in(module_name: &[u8], original_fn: *const (), hook_fn: *const ()) -> usize {
    let module = GetModuleHandleA(module_name.as_ptr());
    if module.is_null() {
        return 0;
    }
    patch_module(module, original_fn, hook_fn)
}

unsafe fn patch_module(
    module: HMODULE,
    original_fn: *const (),
    hook_fn: *const (),
) -> usize {
    if module.is_null() {
        return 0;
    }
    let base = module as usize;
    let dos = base as *const ImageDosHeader;
    if (*dos).e_magic != 0x5A4D {
        return 0;
    }

    let nt = (base + (*dos).e_lfanew as usize) as *const ImageNtHeaders64;

    let import_rva = (*nt).optional_header.data_directory[1].virtual_address as usize;
    if import_rva == 0 {
        return 0;
    }

    let mut desc = (base + import_rva) as *const ImageImportDescriptor;
    let mut count = 0usize;

    while (*desc).name != 0 {
        let ft = (*desc).first_thunk as usize;
        let iat_slot = (base + ft) as *mut usize;
        let mut i = 0usize;

        loop {
            let val = *iat_slot.add(i);
            if val == 0 {
                break;
            }
            if val == original_fn as usize {
                let target = iat_slot.add(i) as *const std::ffi::c_void;
                let mut old: u32 = 0;
                VirtualProtect(target, std::mem::size_of::<usize>(), PAGE_EXECUTE_READWRITE, &mut old);
                *iat_slot.add(i) = hook_fn as usize;
                VirtualProtect(target, std::mem::size_of::<usize>(), old, &mut old);
                count += 1;
            }
            i += 1;
        }

        desc = desc.add(1);
    }

    count
}

/// Resolve the address of an exported function from an already-loaded DLL.
pub unsafe fn resolve(dll: &[u8], fn_name: &[u8]) -> *const () {
    let module = GetModuleHandleA(dll.as_ptr());
    if module.is_null() {
        return std::ptr::null();
    }
    match GetProcAddress(module, fn_name.as_ptr()) {
        Some(f) => f as *const (),
        None => std::ptr::null(),
    }
}