climperor/gsi_watch.py

"""Watch Dota 2 via Game State Integration and recognize the draft automatically.

Usage:
    python gsi_setup.py       # once: install the GSI cfg into Dota 2
    python gsi_watch.py       # then leave this running while you play
    python gsi_watch.py --once            # capture+recognize right now, no GSI
    python gsi_watch.py --port 3223
    python gsi_watch.py --states HERO_SELECTION,STRATEGY_TIME
    python gsi_watch.py --dump-gsi        # force full GSI JSONL logging
    python gsi_watch.py --no-dump-gsi     # disable it

When the game enters hero selection the watcher follows the whole draft,
polling the screen and logging each pick as it is revealed (All Pick reveals
them in waves of 2/2/1 per team). It also works out which slot is you and
which lane role you queued for. The result lands in results/draft_<ts>.json.

If config.json has no calibrated slots yet the watcher runs in capture-only
mode: it still saves frames to samples/raw/<matchid>/ so you can calibrate from them.

With gsi.dump_payloads (default on), every POST body is appended to
samples/raw/<matchid>/gsi.jsonl for later analysis.
"""

import json
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path

import mss

from capture import append_gsi_payload, grab_frame, is_dota_foreground, raw_dir_for_match, save_frame
from common import ROOT, load_config, load_template_library
from draft_session import HERO_SELECTION, DraftSession, describe, gsi_slot, loc
from overlay import DraftOverlay
from recognize import recognize_image
from roles import detect_roles

RESULTS_DIR = ROOT / "results"

# entering any of these means a new match is starting - allow triggering again
RESET_STATES = {
    "DOTA_GAMERULES_STATE_INIT",
    "DOTA_GAMERULES_STATE_WAIT_FOR_PLAYERS_TO_LOAD",
    "DOTA_GAMERULES_STATE_POST_GAME",
    "DOTA_GAMERULES_STATE_DISCONNECT",
}


class Watcher:
    """Turns GSI state changes into capture+recognize runs."""

    def __init__(self, cfg: dict, calibrated: bool, verbose: bool = False):
        self.cfg = cfg
        self.calibrated = calibrated
        self.verbose = verbose
        gsi = cfg.get("gsi", {})
        self.trigger_states = set(gsi.get("trigger_states", ["DOTA_GAMERULES_STATE_STRATEGY_TIME"]))
        self.interval = gsi.get("capture_interval", 1.0)
        self.dump_payloads = bool(gsi.get("dump_payloads", True))

        self.library = load_template_library() if calibrated else []
        self.last_state: str | None = None
        self.last_match_id: str | None = None
        self.handled_matches: set[str] = set()
        self.connected = False
        self.busy = threading.Lock()
        self.dump_lock = threading.Lock()
        self.self_info: dict = {}
        self._dump_announced: set[str] = set()
        # One overlay for the whole process - Tk does not like create/destroy per match.
        self.overlay = None
        if calibrated and bool((cfg.get("overlay") or {}).get("enabled", True)):
            try:
                self.overlay = DraftOverlay(cfg)
                print("[draft] role-tag overlay ready", flush=True)
            except Exception as e:  # noqa: BLE001
                print(f"[draft] overlay disabled: {e}", flush=True)
                self.overlay = None

    def on_payload(self, payload: dict) -> None:
        if not self.connected:
            self.connected = True
            name = (payload.get("provider") or {}).get("name", "Dota 2")
            print(f"[gsi] connected to {name}", flush=True)

        p = payload.get("player") or {}
        h = payload.get("hero") or {}
        self.self_info = {
            "name": p.get("name"),
            "team": p.get("team_name"),
            "team_slot": p.get("team_slot"),
            "hero": (h.get("name") or "").replace("npc_dota_hero_", "") or None,
            "accountid": p.get("accountid"),
            "steamid": p.get("steamid"),
        }

        state = (payload.get("map") or {}).get("game_state")
        match_id = str((payload.get("map") or {}).get("matchid") or "no-match")
        self.last_match_id = match_id
        if self.dump_payloads:
            self._dump_gsi(match_id, payload)
        if state is None:
            # main menu / no active match
            if self.last_state is not None:
                print("[gsi] left match (back in menu)", flush=True)
                self.last_state = None
            return

        if state in RESET_STATES and self.handled_matches:
            self.handled_matches.clear()

        if state != self.last_state:
            m = payload.get("map") or {}
            extras = {k: m[k] for k in ("game_mode", "lobby_type", "customgamename", "name") if k in m}
            print(f"[gsi] {self.last_state} -> {state}  (match {match_id}) {extras or ''}", flush=True)
            if not getattr(self, "_dumped_map_keys", False) and m:
                self._dumped_map_keys = True
                print(f"[gsi] map keys: {sorted(m.keys())}", flush=True)
                print(f"[gsi] player keys: {sorted(p.keys())}", flush=True)
                print(f"[gsi] self: {self.self_info} -> top-bar slot {gsi_slot(self.self_info)}", flush=True)
            self.last_state = state

        if state not in self.trigger_states:
            return
        # one tracking session per match, however far into the draft we joined
        key = f"{match_id}:draft"
        if key in self.handled_matches:
            return
        threading.Thread(target=self.track, args=(match_id, state, key), daemon=True).start()

    def _dump_gsi(self, match_id: str, payload: dict) -> None:
        """Persist the full POST body; announce the path once per match folder."""
        try:
            with self.dump_lock:
                path = append_gsi_payload(match_id, payload)
            if match_id not in self._dump_announced:
                self._dump_announced.add(match_id)
                print(f"[gsi] dumping payloads -> {path}", flush=True)
        except OSError as e:
            print(f"[gsi] dump failed: {e}", flush=True)

    def track(self, match_id: str, state: str, key: str) -> dict | None:
        """Follow the draft from here to the end, recording every reveal."""
        if not self.busy.acquire(blocking=False):
            return None
        try:
            if key in self.handled_matches:
                return None
            self.handled_matches.add(key)
            if not self.calibrated:
                return self._capture_only(match_id, state)

            late = " (joined late)" if state != HERO_SELECTION else ""
            print(f"[draft] tracking match {match_id} from {state}{late}", flush=True)
            session = DraftSession(self.cfg, self.library, overlay=self.overlay)
            summary = session.run(match_id, lambda: self.last_state, lambda: self.self_info)
            if not summary or summary["recognized"] == 0:
                print("[draft] nothing recognized - no result written", flush=True)
                return None
            self.report_session(summary)
            return summary
        finally:
            self.busy.release()

    def _capture_only(self, match_id: str, state: str) -> None:
        """No calibration yet: just bank a few frames to calibrate from later."""
        out = raw_dir_for_match(match_id)
        require_fg = bool(self.cfg.get("gsi", {}).get("require_foreground", True))
        print(f"[run] capture-only -> {out} (state={state})", flush=True)
        saved = 0
        skipped_fg = False
        with mss.MSS() as sct:
            # Retry until 3 usable frames or a short deadline so Alt-Tab does not
            # bank desktop screenshots for calibration.
            deadline = time.monotonic() + max(self.interval * 12, 15.0)
            while saved < 3 and time.monotonic() < deadline:
                if require_fg and not is_dota_foreground():
                    if not skipped_fg:
                        print("[run] Dota 2 not foreground - waiting to capture", flush=True)
                        skipped_fg = True
                    time.sleep(self.interval)
                    continue
                if skipped_fg:
                    print("[run] Dota 2 foreground again - capturing", flush=True)
                    skipped_fg = False
                saved += 1
                path = save_frame(grab_frame(sct), out, prefix="draft")
                print(f"[run] capture-only {saved}/3: {path}", flush=True)
                time.sleep(self.interval)
        if saved == 0:
            print("[run] capture-only got no frames (Dota never foreground)", flush=True)
        else:
            print("[run] no calibrated slots - run calibrate.py on one of these frames", flush=True)
        return None

    def report_session(self, summary: dict) -> None:
        RESULTS_DIR.mkdir(exist_ok=True)
        out_path = RESULTS_DIR / f"draft_{time.strftime('%Y%m%d_%H%M%S')}.json"
        out_path.write_text(json.dumps(summary, ensure_ascii=False, indent=1), encoding="utf-8")
        print("", flush=True)
        for line in describe(summary):
            print(line, flush=True)
        print(f"saved  : {out_path}", flush=True)

    def run_once(self) -> dict | None:
        """One snapshot of whatever is on screen right now, for manual checks."""
        if not is_dota_foreground():
            print("[run] warning: Dota 2 is not the foreground window", flush=True)
        frame = grab_frame()
        out = raw_dir_for_match(self.last_match_id)
        if not self.calibrated:
            print(f"[run] capture-only: {save_frame(frame, out, prefix='draft')}", flush=True)
            return None

        result = recognize_image(frame, self.cfg, self.library)
        found = detect_roles(frame, self.cfg)
        print(f"[run] {result['recognized']}/10 slots, {result['elapsed_ms']}ms", flush=True)
        print(f"radiant: {', '.join(loc(r['hero']) for r in result['radiant'])}", flush=True)
        print(f"dire   : {', '.join(loc(r['hero']) for r in result['dire'])}", flush=True)
        slot = gsi_slot(self.self_info)
        team = self.self_info.get("team") or found.get("self_team")
        if slot:
            role = found["roles"].get(slot)
            pos = f"position {role['position']} ({role['label']})" if role else "position unknown"
            print(f"you    : slot {slot} {team} - {pos}", flush=True)
        elif not slot:
            print("you    : slot unknown (waiting for GSI team_slot)", flush=True)
        result["roles"] = found
        return result


class SingleBindServer(HTTPServer):
    """Fail loudly when the port is taken.

    Windows honours SO_REUSEADDR literally, so the stdlib default would let a
    second watcher bind 3223 silently and steal half the GSI payloads.
    """

    allow_reuse_address = False


def make_handler(watcher: Watcher):
    class Handler(BaseHTTPRequestHandler):
        protocol_version = "HTTP/1.1"


        def do_POST(self):  # noqa: N802 - required by BaseHTTPRequestHandler
            length = int(self.headers.get("Content-Length", 0))
            body = self.rfile.read(length) if length else b"{}"
            self.send_response(200)
            self.send_header("Content-Length", "0")
            self.end_headers()
            try:
                watcher.on_payload(json.loads(body.decode("utf-8")))
            except (ValueError, UnicodeDecodeError) as e:
                print(f"[gsi] bad payload: {e}", flush=True)

        def log_message(self, fmt, *args):
            if watcher.verbose:
                super().log_message(fmt, *args)

    return Handler


def main() -> None:
    # Keep progress visible when stdout is a pipe or file, not just a console,
    # and force UTF-8 so the Chinese hero names survive the default Windows
    # console code page (which mangles them into mojibake).
    for stream in (sys.stdout, sys.stderr):
        stream.reconfigure(encoding="utf-8", errors="replace", line_buffering=True)

    args = sys.argv[1:]
    cfg = load_config()
    calibrated = bool(cfg.get("slots"))
    verbose = "--verbose" in args

    if not calibrated:
        print("WARNING: config.json has no calibrated slots - running in capture-only mode.")
        print("         Play one draft, then: python calibrate.py samples/raw/<matchid>/<frame>.png")
        print()

    watcher = Watcher(cfg, calibrated, verbose)
    if "--dump-gsi" in args:
        watcher.dump_payloads = True
    if "--no-dump-gsi" in args:
        watcher.dump_payloads = False

    if "--states" in args:
        names = args[args.index("--states") + 1].split(",")
        watcher.trigger_states = {
            n if n.startswith("DOTA_GAMERULES_STATE_") else f"DOTA_GAMERULES_STATE_{n}"
            for n in (s.strip().upper() for s in names)
            if n
        }

    if "--once" in args:
        watcher.run_once()
        return

    port = int(args[args.index("--port") + 1]) if "--port" in args else cfg.get("gsi", {}).get("port", 3223)
    try:
        server = SingleBindServer(("127.0.0.1", port), make_handler(watcher))
    except OSError as e:
        sys.exit(
            f"cannot bind 127.0.0.1:{port} ({e}).\n"
            "Another gsi_watch.py is probably still running - stop it first:\n"
            "  Get-CimInstance Win32_Process -Filter \"Name='python.exe'\" |\n"
            "    Where-Object { $_.CommandLine -like '*gsi_watch*' } |\n"
            "    ForEach-Object { Stop-Process -Id $_.ProcessId -Force }"
        )
    print(f"listening on http://127.0.0.1:{port}/  (Ctrl+C to stop)")
    print(f"trigger states: {', '.join(sorted(watcher.trigger_states))}")
    print(f"gsi dump   : {'on -> samples/raw/<matchid>/gsi.jsonl' if watcher.dump_payloads else 'off'}")
    if calibrated:
        print(f"template library: {len(watcher.library)} entries")
    print("waiting for Dota 2 ... (needs -gamestateintegration launch option)")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nstopped.")
    finally:
        if watcher.overlay is not None:
            watcher.overlay.close()
        server.server_close()


if __name__ == "__main__":
    main()