climperor/fetch_stratz.py

"""Fetch STRATZ hero-vs-hero matchups or teammate synergies.

Requires STRATZ_API_TOKEN in the environment or a local .env file
(gitignored). Attribution: data from https://stratz.com

Usage:
  python fetch_stratz.py --mode matchups     # -> data/stratz_matchups.json
  python fetch_stratz.py --mode synergies     # -> data/synergies.json
  python fetch_stratz.py --mode matchups --delay 0.45
"""

from __future__ import annotations

import argparse
import json
import os
import time
import urllib.error
import urllib.request
from datetime import datetime, timezone
from pathlib import Path

from common import DATA, ROOT
from grid import hero_table

API = "https://api.stratz.com/graphql"

# GraphQL edge field ("vs" for matchups, "with" for synergies) and the
# score field name used in the output JSON.
MODE_CONFIG = {
    "matchups": {
        "edge_field": "vs",
        "score_field": "advantage",
        "label": "matchups",
        "default_out": DATA / "stratz_matchups.json",
    },
    "synergies": {
        "edge_field": "with",
        "score_field": "synergy",
        "label": "synergies",
        "default_out": DATA / "synergies.json",
    },
}


def _build_query(edge_field: str) -> str:
    return f"""
query($id: Short!) {{
  heroStats {{
    heroVsHeroMatchup(heroId: $id) {{
      advantage {{
        heroId
        {edge_field} {{
          heroId2
          matchCount
          winCount
          synergy
          winsAverage
        }}
      }}
    }}
  }}
}}
"""


def load_token() -> str:
    env = os.environ.get("STRATZ_API_TOKEN", "").strip()
    if env:
        return env
    path = ROOT / ".env"
    if path.is_file():
        for line in path.read_text(encoding="utf-8").splitlines():
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            k, v = line.split("=", 1)
            if k.strip() == "STRATZ_API_TOKEN":
                return v.strip().strip('"').strip("'")
    raise SystemExit(
        "STRATZ_API_TOKEN missing. Set env or add it to .env (see .env.example)."
    )


def gql(token: str, query: str, variables: dict | None = None) -> dict:
    body = {"query": query}
    if variables:
        body["variables"] = variables
    req = urllib.request.Request(
        API,
        data=json.dumps(body).encode(),
        headers={
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
            "User-Agent": "climperor",
        },
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=60) as resp:
        payload = json.loads(resp.read().decode())
    if payload.get("errors"):
        raise RuntimeError(str(payload["errors"][:2]))
    return payload["data"]


def fetch_one(
    token: str, hero_id: int, *, edge_field: str, score_field: str
) -> dict[str, dict]:
    query = _build_query(edge_field)
    data = gql(token, query, {"id": hero_id})
    block = (((data or {}).get("heroStats") or {}).get("heroVsHeroMatchup") or {})
    rows = block.get("advantage") or []
    cell: dict[str, dict] = {}
    for row in rows:
        for pair in row.get(edge_field) or []:
            other = pair.get("heroId2")
            games = int(pair.get("matchCount") or 0)
            wins = int(pair.get("winCount") or 0)
            if other is None or games <= 0:
                continue
            cell[str(int(other))] = {
                "games": games,
                "wins": wins,
                score_field: float(pair.get("synergy") or 0.0),
                "wr": float(pair.get("winsAverage") or (wins / games)),
            }
    return cell


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--mode", choices=list(MODE_CONFIG), required=True)
    ap.add_argument("--delay", type=float, default=0.45)
    ap.add_argument("--out", type=Path, default=None)
    args = ap.parse_args()

    mc = MODE_CONFIG[args.mode]
    out_path = args.out or mc["default_out"]
    edge_field = mc["edge_field"]
    score_field = mc["score_field"]
    label = mc["label"]

    token = load_token()
    ids = sorted({int(h["id"]) for h in hero_table()})
    by_hero: dict[str, dict] = {}
    if out_path.is_file():
        try:
            prev = json.loads(out_path.read_text(encoding="utf-8"))
            by_hero = dict(prev.get("by_hero") or {})
            print(f"resuming with {len(by_hero)} heroes already cached", flush=True)
        except (OSError, json.JSONDecodeError):
            pass

    pending = [i for i in ids if str(i) not in by_hero]
    print(f"fetching {len(pending)} / {len(ids)} heroes -> {out_path}", flush=True)
    for n, hid in enumerate(pending, start=1):
        try:
            cell = fetch_one(token, hid, edge_field=edge_field, score_field=score_field)
        except (urllib.error.URLError, TimeoutError, RuntimeError, json.JSONDecodeError) as e:
            print(f"  [{n}/{len(pending)}] hero {hid} failed: {e}", flush=True)
            time.sleep(args.delay * 2)
            continue
        by_hero[str(hid)] = cell
        print(f"  [{n}/{len(pending)}] hero {hid}: {len(cell)} {label}", flush=True)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        payload = {
            "fetched_at": datetime.now(timezone.utc).isoformat(),
            "source": "stratz",
            "attribution": "https://stratz.com",
            "by_hero": by_hero,
        }
        out_path.write_text(
            json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
            encoding="utf-8",
        )
        time.sleep(args.delay)

    print(f"done: {len(by_hero)} heroes in {out_path}", flush=True)


if __name__ == "__main__":
    main()