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

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.
"""

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

import mss

from capture import grab_frame, 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 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.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.self_info: dict = {}

    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,
        }

        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 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 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)
            summary = session.run(match_id, lambda: self.last_state, lambda: self.self_info)
            if 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)
        print(f"[run] capture-only -> {out} (state={state})", flush=True)
        with mss.MSS() as sct:
            for attempt in range(1, 4):
                path = save_frame(grab_frame(sct), out, prefix="draft")
                print(f"[run] capture-only {attempt}/3: {path}", flush=True)
                time.sleep(self.interval)
        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."""
        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) or found["self_slot"]
        if slot:
            role = found["roles"].get(slot)
            pos = f"position {role['position']} ({role['label']})" if role else "position unknown"
            print(f"you    : slot {slot} {found['self_team']} - {pos}", 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 "--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))}")
    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:
        server.server_close()


if __name__ == "__main__":
    main()