/// Intercepts ConnectEx (EA/DirtySDK's preferred async connect API).
///
/// DirtySDK calls WSAIoctl(SIO_GET_EXTENSION_FUNCTION_POINTER, WSAID_CONNECTEX) once at
/// startup to get a ConnectEx function pointer, bypassing all IAT hooks.  We hook WSAIoctl
/// inline so that when it returns a ConnectEx pointer we swap it for our own wrapper.
use core::sync::atomic::{AtomicUsize, Ordering};
use core::ffi::c_void;

// Address rewriting (v4 + v6) is shared from connect_hook::redirect_if_ea, so the port
// constants and sockaddr structs no longer live here.

// SIO_GET_EXTENSION_FUNCTION_POINTER
const SIO_GET_EXT_FN: u32 = 0xC8000006;

// WSAID_CONNECTEX = {25A207B9-DDF3-4660-8EE9-76E58C74063E}
const CONNECTEX_GUID: [u8; 16] = [
    0xB9, 0x07, 0xA2, 0x25,
    0xF3, 0xDD, 0x60, 0x46,
    0x8E, 0xE9, 0x76, 0xE5, 0x8C, 0x74, 0x06, 0x3E,
];

// The real ConnectEx pointer, saved after WSAIoctl returns it
static REAL_CONNECTEX: AtomicUsize = AtomicUsize::new(0);

// ConnectEx function signature
type ConnectExFn = unsafe extern "system" fn(
    s: usize,
    name: *const u8,
    namelen: i32,
    send_buf: *const c_void,
    send_data_len: u32,
    bytes_sent: *mut u32,
    overlapped: *mut c_void,
) -> i32;

// WSAIoctl function address (for inline unhook/rehook)
static WSAIOCTL_ADDR: AtomicUsize = AtomicUsize::new(0);
static mut WSAIOCTL_ORIG: [u8; 14] = [0u8; 14];

type WsaIoctlFn = unsafe extern "system" fn(
    s: usize,
    code: u32,
    in_buf: *const c_void,
    in_len: u32,
    out_buf: *mut c_void,
    out_len: u32,
    bytes_ret: *mut u32,
    overlapped: *mut c_void,
    completion: *const c_void,
) -> i32;

unsafe fn write_hook(target: *mut u8, dest: u64) {
    use windows_sys::Win32::System::Memory::{VirtualProtect, PAGE_EXECUTE_READWRITE};
    let mut old: u32 = 0;
    VirtualProtect(target as _, 14, PAGE_EXECUTE_READWRITE, &mut old);
    target.write(0xFF); target.add(1).write(0x25);
    (target.add(2) as *mut u32).write(0u32);
    (target.add(6) as *mut u64).write(dest);
    VirtualProtect(target as _, 14, old, &mut old);
}

unsafe fn restore_wsaioctl(target: *mut u8) {
    use windows_sys::Win32::System::Memory::{VirtualProtect, PAGE_EXECUTE_READWRITE};
    let mut old: u32 = 0;
    VirtualProtect(target as _, 14, PAGE_EXECUTE_READWRITE, &mut old);
    core::ptr::copy_nonoverlapping(WSAIOCTL_ORIG.as_ptr(), target, 14);
    VirtualProtect(target as _, 14, old, &mut old);
}

/// Our ConnectEx wrapper: redirects EA ports to 127.0.0.1
unsafe extern "system" fn hooked_connectex(
    s: usize,
    name: *const u8,
    namelen: i32,
    send_buf: *const c_void,
    send_data_len: u32,
    bytes_sent: *mut u32,
    overlapped: *mut c_void,
) -> i32 {
    let real_fn: ConnectExFn = core::mem::transmute(REAL_CONNECTEX.load(Ordering::Relaxed));

    // Milestone-0 transport watch (self-gates on OPENFUT_TRANSPORT_WATCH).
    crate::transport_watch::note_connect("ConnectEx", name, namelen, s);

    // Share the one redirect implementation (v4 + v6) with connect_hook, so ConnectEx
    // dials get the same IPv6 handling as plain connect().
    if let Some((buf, len)) = crate::connect_hook::redirect_if_ea(name, namelen) {
        return real_fn(s, buf.as_ptr(), len, send_buf, send_data_len, bytes_sent, overlapped);
    }
    real_fn(s, name, namelen, send_buf, send_data_len, bytes_sent, overlapped)
}

/// Our WSAIoctl hook: when ConnectEx is requested, save the real pointer and return ours
pub unsafe extern "system" fn hooked_wsaioctl(
    s: usize,
    code: u32,
    in_buf: *const c_void,
    in_len: u32,
    out_buf: *mut c_void,
    out_len: u32,
    bytes_ret: *mut u32,
    overlapped: *mut c_void,
    completion: *const c_void,
) -> i32 {
    let addr = WSAIOCTL_ADDR.load(Ordering::Relaxed) as *mut u8;

    // Call the real WSAIoctl via unhook/rehook
    restore_wsaioctl(addr);
    let result = {
        let f: WsaIoctlFn = core::mem::transmute(addr);
        f(s, code, in_buf, in_len, out_buf, out_len, bytes_ret, overlapped, completion)
    };
    write_hook(addr, hooked_wsaioctl as u64);

    // If this was a ConnectEx request that succeeded, swap the pointer
    if result == 0
        && code == SIO_GET_EXT_FN
        && in_len == 16
        && !in_buf.is_null()
    {
        let guid = core::slice::from_raw_parts(in_buf as *const u8, 16);
        if guid == CONNECTEX_GUID
            && out_len >= 8
            && !out_buf.is_null()
        {
            let out_ptr = out_buf as *mut usize;
            let real_addr = *out_ptr;
            if REAL_CONNECTEX.compare_exchange(0, real_addr, Ordering::Relaxed, Ordering::Relaxed).is_ok() {
                crate::write_log(&format!("connectex_hook: intercepted ConnectEx @ {real_addr:#x}\n"));
            }
            // Return our hook instead
            *out_ptr = hooked_connectex as usize;
        }
    }
    result
}

pub unsafe fn install_wsaioctl_hook() -> bool {
    use windows_sys::Win32::System::LibraryLoader::{GetModuleHandleA, GetProcAddress};
    let ws2 = GetModuleHandleA(b"ws2_32.dll\0".as_ptr());
    if ws2.is_null() { return false; }
    let fn_ptr = match GetProcAddress(ws2, b"WSAIoctl\0".as_ptr()) {
        Some(f) => f as *mut u8,
        None => return false,
    };
    core::ptr::copy_nonoverlapping(fn_ptr, WSAIOCTL_ORIG.as_mut_ptr(), 14);
    WSAIOCTL_ADDR.store(fn_ptr as usize, Ordering::Relaxed);
    write_hook(fn_ptr, hooked_wsaioctl as u64);
    true
}