asdf-overlay 1.2.3

use core::ffi::c_void;

use asdf_overlay_hook::DetourHook;
use once_cell::sync::OnceCell;
use tracing::debug;
use windows::{
    Win32::{
        Foundation::{HWND, POINT, RECT},
        UI::{
            Input::{
                HRAWINPUT, KeyboardAndMouse::GetActiveWindow, RAW_INPUT_DATA_COMMAND_FLAGS,
                RAWINPUT, RAWINPUTHEADER, RID_HEADER, RID_INPUT,
            },
            WindowsAndMessaging::GetForegroundWindow,
        },
    },
    core::BOOL,
};

use crate::{
    backend::{Backends, window::InputBlockData},
    hook::proc::message_reading,
};

windows::core::link!("user32.dll" "system" fn ClipCursor(lprect: *const RECT) -> BOOL);
windows::core::link!("user32.dll" "system" fn SetCursorPos(x: i32, y: i32) -> BOOL);

windows::core::link!("user32.dll" "system" fn GetClipCursor(lprect: *mut RECT) -> BOOL);
windows::core::link!("user32.dll" "system" fn GetCursorPos(lppoint: *mut POINT) -> BOOL);
windows::core::link!("user32.dll" "system" fn GetPhysicalCursorPos(lppoint: *mut POINT) -> BOOL);
windows::core::link!("user32.dll" "system" fn GetKeyboardState(buf: *mut u8) -> BOOL);
windows::core::link!("user32.dll" "system" fn GetKeyState(vkey: i32) -> i16);
windows::core::link!("user32.dll" "system" fn GetAsyncKeyState(vkey: i32) -> i16);
windows::core::link!(
    "user32.dll" "system"
    fn GetRawInputData(
        hrawinput: HRAWINPUT,
        uicommand: RAW_INPUT_DATA_COMMAND_FLAGS,
        pdata: *mut c_void,
        pcbsize: *mut u32,
        cbsizeheader: u32,
    ) -> u32
);
windows::core::link!("user32.dll" "system" fn GetRawInputBuffer(pdata: *mut RAWINPUT, pcbsize: *mut u32, cbsizeheader: u32) -> u32);

struct Hook {
    clip_cursor: DetourHook<ClipCursorFn>,
    set_cursor_pos: DetourHook<SetCursorFn>,

    get_clip_cursor: DetourHook<GetClipCursorFn>,
    get_cursor_pos: DetourHook<GetCursorPos>,
    get_physical_cursor_pos: DetourHook<GetPhysicalCursorPos>,
    get_async_key_state: DetourHook<GetAsyncKeyStateFn>,
    get_key_state: DetourHook<GetKeyStateFn>,
    get_keyboard_state: DetourHook<GetKeyboardStateFn>,
    get_raw_input_data: DetourHook<GetRawInputDataFn>,
    get_raw_input_buffer: DetourHook<GetRawInputBufferFn>,
}
static HOOK: OnceCell<Hook> = OnceCell::new();

type ClipCursorFn = unsafe extern "system" fn(*const RECT) -> BOOL;
type SetCursorFn = unsafe extern "system" fn(i32, i32) -> BOOL;

type GetClipCursorFn = unsafe extern "system" fn(*mut RECT) -> BOOL;
type GetCursorPos = unsafe extern "system" fn(*mut POINT) -> BOOL;
type GetPhysicalCursorPos = unsafe extern "system" fn(*mut POINT) -> BOOL;
type GetAsyncKeyStateFn = unsafe extern "system" fn(i32) -> i16;
type GetKeyStateFn = unsafe extern "system" fn(i32) -> i16;
type GetKeyboardStateFn = unsafe extern "system" fn(*mut u8) -> BOOL;
type GetRawInputDataFn = unsafe extern "system" fn(
    HRAWINPUT,
    RAW_INPUT_DATA_COMMAND_FLAGS,
    *mut c_void,
    *mut u32,
    u32,
) -> u32;
type GetRawInputBufferFn = unsafe extern "system" fn(*mut RAWINPUT, *mut u32, u32) -> u32;

pub fn hook() -> anyhow::Result<()> {
    HOOK.get_or_try_init(|| unsafe {
        debug!("hooking ClipCursor");
        let clip_cursor = DetourHook::attach(ClipCursor as _, hooked_clip_cursor as _)?;

        debug!("hooking SetCursorPos");
        let set_cursor_pos = DetourHook::attach(SetCursorPos as _, hooked_set_cursor_pos as _)?;

        debug!("hooking GetClipCursor");
        let get_clip_cursor = DetourHook::attach(GetClipCursor as _, hooked_get_clip_cursor as _)?;

        debug!("hooking GetCursorPos");
        let get_cursor_pos = DetourHook::attach(GetCursorPos as _, hooked_get_cursor_pos as _)?;

        debug!("hooking GetPhysicalCursorPos");
        let get_physical_cursor_pos = DetourHook::attach(
            GetPhysicalCursorPos as _,
            hooked_get_physical_cursor_pos as _,
        )?;

        debug!("hooking GetAsyncKeyState");
        let get_async_key_state =
            DetourHook::attach(GetAsyncKeyState as _, hooked_get_async_key_state as _)?;

        debug!("hooking GetKeyState");
        let get_key_state = DetourHook::attach(GetKeyState as _, hooked_get_key_state as _)?;

        debug!("hooking GetKeyboardState");
        let get_keyboard_state =
            DetourHook::attach(GetKeyboardState as _, hooked_get_keyboard_state as _)?;

        debug!("hooking GetRawInputData");
        let get_raw_input_data =
            DetourHook::attach(GetRawInputData as _, hooked_get_raw_input_data as _)?;

        debug!("hooking GetRawInputBuffer");
        let get_raw_input_buffer =
            DetourHook::attach(GetRawInputBuffer as _, hooked_get_raw_input_buffer as _)?;

        Ok::<_, anyhow::Error>(Hook {
            clip_cursor,
            set_cursor_pos,

            get_clip_cursor,
            get_cursor_pos,
            get_physical_cursor_pos,
            get_async_key_state,
            get_key_state,
            get_keyboard_state,
            get_raw_input_data,
            get_raw_input_buffer,
        })
    })?;

    Ok(())
}

#[inline]
fn active_hwnd_can_block() -> Option<HWND> {
    let hwnd = unsafe { GetActiveWindow() };

    if !hwnd.is_invalid() && !message_reading() {
        Some(hwnd)
    } else {
        None
    }
}

#[inline]
fn active_hwnd_with<R>(f: impl FnOnce(&mut InputBlockData) -> R) -> Option<R> {
    Backends::with_backend(active_hwnd_can_block()?.0 as _, |backend| {
        let mut proc = backend.proc.lock();
        Some(f(proc.blocking_state.as_mut()?))
    })
    .flatten()
}

#[inline]
fn foreground_hwnd_input_blocked() -> bool {
    let hwnd = unsafe { GetForegroundWindow() };

    !hwnd.is_invalid()
        && Backends::with_backend(hwnd.0 as _, |backend| backend.proc.lock().input_blocking())
            .unwrap_or(false)
}

#[inline]
fn active_hwnd_input_blocked() -> bool {
    active_hwnd_with(|_| ()).is_some()
}

#[tracing::instrument]
extern "system" fn hooked_clip_cursor(lprect: *const RECT) -> BOOL {
    if active_hwnd_with(|data| {
        data.clip_cursor = unsafe { lprect.as_ref() }.copied();
    })
    .is_some()
    {
        return BOOL(1);
    }

    unsafe { HOOK.wait().clip_cursor.original_fn()(lprect) }
}

#[tracing::instrument]
extern "system" fn hooked_set_cursor_pos(x: i32, y: i32) -> BOOL {
    if foreground_hwnd_input_blocked() {
        return BOOL(1);
    }

    unsafe { HOOK.wait().set_cursor_pos.original_fn()(x, y) }
}

#[tracing::instrument]
extern "system" fn hooked_get_clip_cursor(lprect: *mut RECT) -> BOOL {
    match active_hwnd_with(|data| data.clip_cursor).flatten() {
        Some(rect) => {
            unsafe { lprect.write(rect) };
            BOOL(1)
        }
        None => unsafe { HOOK.wait().get_clip_cursor.original_fn()(lprect) },
    }
}

#[tracing::instrument]
extern "system" fn hooked_get_cursor_pos(lppoint: *mut POINT) -> BOOL {
    if foreground_hwnd_input_blocked() {
        // Return a fixed position instead of the real cursor position to prevent games from tracking mouse movement
        if !lppoint.is_null() {
            unsafe {
                lppoint.write(POINT { x: 0, y: 0 });
            }
        }
        return BOOL(1);
    }

    unsafe { HOOK.wait().get_cursor_pos.original_fn()(lppoint) }
}

#[tracing::instrument]
extern "system" fn hooked_get_physical_cursor_pos(lppoint: *mut POINT) -> BOOL {
    if foreground_hwnd_input_blocked() {
        // Return a fixed position instead of the real cursor position to prevent games from tracking mouse movement
        if !lppoint.is_null() {
            unsafe {
                lppoint.write(POINT { x: 0, y: 0 });
            }
        }
        return BOOL(1);
    }

    unsafe { HOOK.wait().get_physical_cursor_pos.original_fn()(lppoint) }
}

#[tracing::instrument]
extern "system" fn hooked_get_async_key_state(vkey: i32) -> i16 {
    if foreground_hwnd_input_blocked() {
        return 0;
    }

    unsafe { HOOK.wait().get_async_key_state.original_fn()(vkey) }
}

#[tracing::instrument]
extern "system" fn hooked_get_key_state(vkey: i32) -> i16 {
    if active_hwnd_input_blocked() {
        return 0;
    }

    unsafe { HOOK.wait().get_key_state.original_fn()(vkey) }
}

#[tracing::instrument]
extern "system" fn hooked_get_keyboard_state(buf: *mut u8) -> BOOL {
    if active_hwnd_input_blocked() {
        // buf is 256 bytes array according to doc.
        unsafe {
            buf.write_bytes(0u8, 256);
        };
        return BOOL(1);
    }

    unsafe { HOOK.wait().get_keyboard_state.original_fn()(buf) }
}

#[tracing::instrument]
extern "system" fn hooked_get_raw_input_data(
    hrawinput: HRAWINPUT,
    uicommand: RAW_INPUT_DATA_COMMAND_FLAGS,
    pdata: *mut c_void,
    pcbsize: *mut u32,
    cbsizeheader: u32,
) -> u32 {
    if foreground_hwnd_input_blocked() {
        // Determine the expected data size based on the command, matching the
        // real Win32 API behaviour so callers (e.g. Godot 4) that validate the
        // returned size against the queried size do not crash.
        let data_size: u32 = match uicommand {
            RID_HEADER => core::mem::size_of::<RAWINPUTHEADER>() as u32,
            RID_INPUT => core::mem::size_of::<RAWINPUT>() as u32,
            _ => 0,
        };

        if !pcbsize.is_null() {
            unsafe { pcbsize.write(data_size) };
        }

        // Size query: return 0 (success) after writing the required buffer size.
        if pdata.is_null() {
            return 0;
        }

        // Data query: write a dummy HID struct and return the number of bytes
        // written, exactly as the real API would on success.
        // Use RIM_TYPEHID so games treat this as an unknown HID device and
        // ignore the empty payload instead of interpreting zeroed fields as
        // mouse/keyboard input.

        let hid_header = RAWINPUTHEADER {
            dwType: 2, // RIM_TYPEHID
            dwSize: data_size,
            ..Default::default()
        };

        match uicommand {
            RID_HEADER => unsafe {
                pdata.cast::<RAWINPUTHEADER>().write(hid_header);
            },
            RID_INPUT => unsafe {
                pdata.cast::<RAWINPUT>().write(RAWINPUT {
                    header: hid_header,
                    ..Default::default()
                });
            },
            _ => {}
        }

        return data_size;
    }

    unsafe {
        HOOK.wait().get_raw_input_data.original_fn()(
            hrawinput,
            uicommand,
            pdata,
            pcbsize,
            cbsizeheader,
        )
    }
}

#[tracing::instrument]
extern "system" fn hooked_get_raw_input_buffer(
    pdata: *mut RAWINPUT,
    pcbsize: *mut u32,
    cbsizeheader: u32,
) -> u32 {
    if foreground_hwnd_input_blocked() {
        unsafe { *pcbsize = 0 };
        return 0;
    }

    unsafe { HOOK.wait().get_raw_input_buffer.original_fn()(pdata, pcbsize, cbsizeheader) }
}