climperor/pc/capture.py

"""Screen capture helper.

Usage:
    python capture.py                     # single full-screen shot -> samples/raw/
    python capture.py --loop 300 2       # capture every 2s for 300s (Ctrl+C to stop early)

Notes:
    - Dota 2 must run in borderless window or windowed mode; exclusive
      fullscreen may capture a black frame with GDI-based grabbers.
    - Frames are saved as PNG at native resolution, named cap_HHMMSS.png.
"""

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import json
import time

import cv2
import mss
import numpy as np

RAW_DIR = Path(__file__).parent / "samples" / "raw"
GSI_JSONL = "gsi.jsonl"
DOTA_EXE = "dota2.exe"


def is_dota_foreground() -> bool:
    """True when the foreground process is dota2.exe.

    Uses Win32 GetForegroundWindow + QueryFullProcessImageNameW. On non-Windows
    or if the probe fails, returns True so recognition is not blocked.
    """
    if sys.platform != "win32":
        return True
    try:
        import ctypes
        from ctypes import wintypes

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

        hwnd = user32.GetForegroundWindow()
        if not hwnd:
            return False

        pid = wintypes.DWORD()
        user32.GetWindowThreadProcessId(hwnd, ctypes.byref(pid))
        if not pid.value:
            return False

        PROCESS_QUERY_LIMITED_INFORMATION = 0x1000
        handle = kernel32.OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, False, pid.value)
        if not handle:
            return False
        try:
            buf = ctypes.create_unicode_buffer(32768)
            size = wintypes.DWORD(len(buf))
            # QueryFullProcessImageNameW(hProcess, dwFlags, lpExeName, lpdwSize)
            QueryFullProcessImageNameW = kernel32.QueryFullProcessImageNameW
            QueryFullProcessImageNameW.argtypes = [
                wintypes.HANDLE, wintypes.DWORD, wintypes.LPWSTR, ctypes.POINTER(wintypes.DWORD)
            ]
            QueryFullProcessImageNameW.restype = wintypes.BOOL
            if not QueryFullProcessImageNameW(handle, 0, buf, ctypes.byref(size)):
                return False
            return Path(buf.value).name.lower() == DOTA_EXE
        finally:
            kernel32.CloseHandle(handle)
    except Exception:
        return True


def raw_dir_for_match(match_id: str | None = None) -> Path:
    """Per-match screenshot folder under samples/raw/{match_id}/.

    Manual capture (no match) still uses samples/raw/ itself.
    """
    if not match_id:
        return RAW_DIR
    safe = "".join(c for c in str(match_id) if c.isalnum() or c in "-_") or "no-match"
    return RAW_DIR / safe


def append_gsi_payload(match_id: str | None, payload: dict) -> Path:
    """Append one full GSI POST body to samples/raw/{match_id}/gsi.jsonl.

    Each line is {"t": <unix seconds>, "payload": <original body>}.
    """
    out = raw_dir_for_match(match_id)
    out.mkdir(parents=True, exist_ok=True)
    path = out / GSI_JSONL
    record = {"t": time.time(), "payload": payload}
    with path.open("a", encoding="utf-8") as f:
        f.write(json.dumps(record, ensure_ascii=False) + "\n")
    return path


def grab_frame(sct=None) -> np.ndarray:
    """Grab the primary monitor as a BGR image."""
    if sct is None:
        with mss.MSS() as own:
            return grab_frame(own)
    shot = sct.grab(sct.monitors[1])
    return cv2.cvtColor(np.asarray(shot), cv2.COLOR_BGRA2BGR)


def save_frame(img: np.ndarray, out_dir: Path = RAW_DIR, prefix: str = "cap") -> str:
    out_dir.mkdir(parents=True, exist_ok=True)
    stamp = time.strftime("%H%M%S")
    path = out_dir / f"{prefix}_{stamp}.png"
    n = 1
    while path.exists():
        path = out_dir / f"{prefix}_{stamp}_{n}.png"
        n += 1
    cv2.imwrite(str(path), img)
    return str(path)


def main() -> None:
    with mss.MSS() as sct:
        if "--loop" in sys.argv:
            i = sys.argv.index("--loop")
            duration = float(sys.argv[i + 1])
            interval = float(sys.argv[i + 2]) if len(sys.argv) > i + 2 else 2.0
            end = time.time() + duration
            n = 0
            print(f"capturing every {interval}s for {duration}s -> {RAW_DIR}")
            try:
                while time.time() < end:
                    path = save_frame(grab_frame(sct))
                    n += 1
                    print(f"[{n}] {path}")
                    time.sleep(interval)
            except KeyboardInterrupt:
                pass
            print(f"done: {n} frames")
        else:
            print(save_frame(grab_frame(sct)))


if __name__ == "__main__":
    main()