climperor/shared/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 sys
from pathlib import Path

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

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

from shared.grid import hero_table
from shared.paths import ROOT, SHARED_DATA

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": SHARED_DATA / "stratz_matchups.json",
    },
    "synergies": {
        "edge_field": "with",
        "score_field": "synergy",
        "label": "synergies",
        "default_out": SHARED_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:
    for key in (
        "STRATZ_API_TOKEN",
        "KEYZOO_ASSET_API_KEY",
        "KEYZOO_ASSET_SECRET_API_KEY",
        "KEYZOO_ASSET_TOKEN",
    ):
        env = os.environ.get(key, "").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 / keyzoo inject or add 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()