"""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 json
import sys
import time
from pathlib import Path

import cv2
import mss
import numpy as np

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


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()