"""Low-level Windows keyboard hook to block escape combos while ICEYOU is locked.

Blocks (when active):
  - WIN keys (left/right) -> Start menu / Win+anything
  - CTRL+ESC               -> Start menu
  - ALT+TAB, ALT+ESC       -> task switching
  - ALT+F4                 -> close window
  - ALT+SPACE              -> window menu (Restore/Move/Close)
  - CTRL+SHIFT+ESC         -> Task Manager hotkey

CANNOT block (by Windows design):
  - CTRL+ALT+DEL (Secure Attention Sequence)
  - Hard reset / power button
  These require Group Policy or admin / kernel privilege.
"""

import ctypes
import threading
from ctypes import wintypes


user32 = ctypes.WinDLL("user32", use_last_error=True)
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)

WH_KEYBOARD_LL = 13
WM_KEYDOWN = 0x0100
WM_SYSKEYDOWN = 0x0104
WM_QUIT = 0x0012

VK_LWIN = 0x5B
VK_RWIN = 0x5C
VK_TAB = 0x09
VK_ESCAPE = 0x1B
VK_F4 = 0x73
VK_SPACE = 0x20
VK_MENU = 0x12      # ALT
VK_CONTROL = 0x11
VK_SHIFT = 0x10

LRESULT = ctypes.c_long
ULONG_PTR = ctypes.c_size_t


class KBDLLHOOKSTRUCT(ctypes.Structure):
    _fields_ = [
        ("vkCode", wintypes.DWORD),
        ("scanCode", wintypes.DWORD),
        ("flags", wintypes.DWORD),
        ("time", wintypes.DWORD),
        ("dwExtraInfo", ULONG_PTR),
    ]


HOOKPROC = ctypes.WINFUNCTYPE(LRESULT, ctypes.c_int, wintypes.WPARAM, wintypes.LPARAM)

user32.SetWindowsHookExW.restype = wintypes.HHOOK
user32.SetWindowsHookExW.argtypes = [ctypes.c_int, HOOKPROC, wintypes.HINSTANCE, wintypes.DWORD]
user32.UnhookWindowsHookEx.restype = wintypes.BOOL
user32.UnhookWindowsHookEx.argtypes = [wintypes.HHOOK]
user32.CallNextHookEx.restype = LRESULT
user32.CallNextHookEx.argtypes = [wintypes.HHOOK, ctypes.c_int, wintypes.WPARAM, wintypes.LPARAM]
user32.GetMessageW.argtypes = [ctypes.POINTER(wintypes.MSG), wintypes.HWND, wintypes.UINT, wintypes.UINT]
user32.GetMessageW.restype = ctypes.c_int
user32.TranslateMessage.argtypes = [ctypes.POINTER(wintypes.MSG)]
user32.DispatchMessageW.argtypes = [ctypes.POINTER(wintypes.MSG)]
user32.PostThreadMessageW.argtypes = [wintypes.DWORD, wintypes.UINT, wintypes.WPARAM, wintypes.LPARAM]
user32.PostThreadMessageW.restype = wintypes.BOOL
user32.GetAsyncKeyState.argtypes = [ctypes.c_int]
user32.GetAsyncKeyState.restype = ctypes.c_short
kernel32.GetCurrentThreadId.restype = wintypes.DWORD


class KeyboardBlocker:
    """Installs a WH_KEYBOARD_LL hook on its own thread with a message pump."""

    def __init__(self):
        self._hook = None
        self._thread: threading.Thread = None
        self._thread_id: int = 0
        self._active = False
        self._hook_proc = None  # keep ref so GC doesn't free callback
        self._ready = threading.Event()

    @property
    def active(self) -> bool:
        return self._active

    def start(self) -> None:
        if self._active:
            return
        self._ready.clear()
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()
        self._ready.wait(timeout=2.0)

    def stop(self) -> None:
        if not self._active:
            return
        self._active = False
        if self._thread_id:
            user32.PostThreadMessageW(self._thread_id, WM_QUIT, 0, 0)
        if self._thread:
            self._thread.join(timeout=2.0)
        self._thread = None
        self._thread_id = 0

    @staticmethod
    def _is_pressed(vk: int) -> bool:
        return (user32.GetAsyncKeyState(vk) & 0x8000) != 0

    def _hook_callback(self, nCode, wParam, lParam):
        try:
            if nCode == 0 and wParam in (WM_KEYDOWN, WM_SYSKEYDOWN):
                kb = ctypes.cast(lParam, ctypes.POINTER(KBDLLHOOKSTRUCT))[0]
                vk = kb.vkCode

                if vk in (VK_LWIN, VK_RWIN):
                    return 1

                alt = self._is_pressed(VK_MENU)
                ctrl = self._is_pressed(VK_CONTROL)

                if alt and vk in (VK_TAB, VK_ESCAPE, VK_F4, VK_SPACE):
                    return 1
                if ctrl and vk == VK_ESCAPE:
                    return 1
        except Exception:
            pass
        return user32.CallNextHookEx(self._hook, nCode, wParam, lParam)

    def _run(self) -> None:
        self._thread_id = kernel32.GetCurrentThreadId()
        self._hook_proc = HOOKPROC(self._hook_callback)
        self._hook = user32.SetWindowsHookExW(WH_KEYBOARD_LL, self._hook_proc, None, 0)
        if not self._hook:
            err = ctypes.get_last_error()
            print(f"[KeyboardBlocker] SetWindowsHookEx failed (error {err})")
            self._active = False
            self._ready.set()
            return
        self._active = True
        self._ready.set()

        msg = wintypes.MSG()
        try:
            while True:
                rc = user32.GetMessageW(ctypes.byref(msg), None, 0, 0)
                if rc <= 0:  # WM_QUIT (0) or error (-1)
                    break
                user32.TranslateMessage(ctypes.byref(msg))
                user32.DispatchMessageW(ctypes.byref(msg))
        finally:
            if self._hook:
                user32.UnhookWindowsHookEx(self._hook)
                self._hook = None
            self._active = False