climperor/web/fetch_ability_videos.py

"""Download official ability demo clips from dota2.com / Steam CDN.

Source pattern (not GIF):
  https://cdn.steamstatic.com/apps/dota2/videos/dota_react/abilities/{hero}/{file}.webm
  https://cdn.steamstatic.com/apps/dota2/videos/dota_react/abilities/{hero}/{file}.mp4
  https://cdn.steamstatic.com/apps/dota2/videos/dota_react/abilities/{hero}/{file}.jpg

{file} is normally the ability key (e.g. juggernaut_blade_fury). However,
abilities GRANTED by Aghanim's Scepter / Shard use the hero-prefixed upgrade
name on the CDN, not the ability key:
  <hero>_aghanims_scepter   (e.g. juggernaut_swift_slash -> juggernaut_aghanims_scepter)
  <hero>_aghanims_shard
The local file is always saved as {ability_key}.{ext} so the web frontend
(app.js) can resolve it by ability key without knowing the grant type.

Ability keys come from data/hero_abilities.json (run fetch_hero_abilities.py first).
Many innate / shard / facet abilities have no clip (CDN 404) — those are skipped.

Rate limiting: polite delay + jitter between requests; longer backoff on 429/403.

Usage:
  python fetch_ability_videos.py --heroes juggernaut
  python fetch_ability_videos.py
  python fetch_ability_videos.py --delay 2 --jitter 1
  python fetch_ability_videos.py --fmt webm --poster
  python fetch_ability_videos.py --force
"""

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 random
import time
import urllib.error
import urllib.request
from datetime import datetime, timezone

from shared.paths import ABILITY_VIDEOS, DATA

HERO_ABILITIES = DATA / "hero_abilities.json"
CDN_BASE = "https://cdn.steamstatic.com/apps/dota2/videos/dota_react/abilities"
UA = "climperor-ability-video-fetch/1.0 (+https://github.com/local/climperor; respectful crawl)"
REFERER = "https://www.dota2.com/"
MANIFEST = ABILITY_VIDEOS / "manifest.json"


def polite_sleep(delay: float, jitter: float) -> None:
    wait = max(0.0, delay) + random.uniform(0.0, max(0.0, jitter))
    if wait > 0:
        time.sleep(wait)


def http_get(url: str, *, timeout: float = 120) -> tuple[str, bytes | None, int | None]:
    """Return (status, body, http_code). status: ok|missing|rate_limited|error."""
    req = urllib.request.Request(
        url,
        headers={
            "User-Agent": UA,
            "Referer": REFERER,
            "Accept": "*/*",
        },
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            data = resp.read()
            ctype = (resp.headers.get("Content-Type") or "").lower()
            if "text/html" in ctype and len(data) < 4096:
                return "error", None, getattr(resp, "status", 200)
            return "ok", data, getattr(resp, "status", 200)
    except urllib.error.HTTPError as e:
        if e.code == 404:
            return "missing", None, 404
        if e.code in (403, 429):
            return "rate_limited", None, e.code
        return "error", None, e.code
    except Exception:  # noqa: BLE001
        return "error", None, None


def _cdn_name(hero_key: str, ab: dict) -> str:
    """CDN demo-clip filename (without extension) for an ability.

    Scepter/shard-granted abilities are published on the CDN under
    ``<hero>_aghanims_scepter`` / ``<hero>_aghanims_shard``, not the
    ability key. The local file still uses the ability key (see app.js).
    """
    if ab.get("granted_by_scepter"):
        return f"{hero_key}_aghanims_scepter"
    if ab.get("granted_by_shard"):
        return f"{hero_key}_aghanims_shard"
    return str(ab.get("key") or "").strip()


def load_ability_index(*, include_innate: bool = False) -> dict[str, list[dict]]:
    """hero_key -> [{key, cdn_name}] for abilities that may have a demo clip.

    Innates are skipped by default (they almost never have a CDN clip);
    pass include_innate=True to probe them too.
    """
    if not HERO_ABILITIES.is_file():
        raise SystemExit(
            f"missing {HERO_ABILITIES}; run: python fetch_hero_abilities.py"
        )
    payload = json.loads(HERO_ABILITIES.read_text(encoding="utf-8"))
    out: dict[str, list[dict]] = {}
    for hero_key, cell in (payload.get("by_hero") or {}).items():
        if not isinstance(cell, dict):
            continue
        entries: list[dict] = []
        for ab in cell.get("abilities") or []:
            if not isinstance(ab, dict):
                continue
            if not include_innate and ab.get("is_innate"):
                continue
            key = str(ab.get("key") or "").strip()
            if not key:
                continue
            entries.append({"key": key, "cdn_name": _cdn_name(hero_key, ab)})
        if entries:
            out[str(hero_key)] = entries
    return out


def load_manifest() -> dict:
    if MANIFEST.is_file():
        try:
            return json.loads(MANIFEST.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            pass
    return {"clips": {}, "missing": [], "meta": {}}


def save_manifest(manifest: dict) -> None:
    ABILITY_VIDEOS.mkdir(parents=True, exist_ok=True)
    manifest["meta"] = {
        "source": "steamcdn/dota_react/abilities",
        "attribution": "https://www.dota2.com",
        "updated_at": datetime.now(timezone.utc).isoformat(),
        "clips": len(manifest.get("clips") or {}),
        "missing": len(manifest.get("missing") or []),
    }
    MANIFEST.write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )


def fetch_one(
    url: str,
    dest: Path,
    *,
    delay: float,
    jitter: float,
    force: bool,
    backoff: float,
) -> str:
    """Download one URL. Returns: ok|skip|missing|fail|rate_limited."""
    if dest.is_file() and dest.stat().st_size > 0 and not force:
        return "skip"

    status, body, code = http_get(url)
    if status == "ok" and body:
        dest.parent.mkdir(parents=True, exist_ok=True)
        tmp = dest.with_suffix(dest.suffix + ".part")
        tmp.write_bytes(body)
        tmp.replace(dest)
        polite_sleep(delay, jitter)
        return "ok"
    if status == "missing":
        polite_sleep(delay * 0.5, jitter * 0.5)
        return "missing"
    if status == "rate_limited":
        print(f"  rate limited ({code}); sleeping {backoff:.0f}s...", flush=True)
        time.sleep(backoff)
        status2, body2, code2 = http_get(url)
        if status2 == "ok" and body2:
            dest.parent.mkdir(parents=True, exist_ok=True)
            tmp = dest.with_suffix(dest.suffix + ".part")
            tmp.write_bytes(body2)
            tmp.replace(dest)
            polite_sleep(delay, jitter)
            return "ok"
        if status2 == "missing":
            polite_sleep(delay * 0.5, jitter * 0.5)
            return "missing"
        print(f"  still blocked ({code2}); aborting batch", flush=True)
        return "rate_limited"
    polite_sleep(delay, jitter)
    return "fail"


def main() -> None:
    ap = argparse.ArgumentParser(
        description="Download dota2.com ability demo videos (rate-limited)"
    )
    ap.add_argument(
        "--heroes",
        default="",
        help="comma-separated hero keys (default: all in hero_abilities.json)",
    )
    ap.add_argument(
        "--fmt",
        choices=("webm", "mp4", "both"),
        default="webm",
        help="video container (default webm; smaller than mp4)",
    )
    ap.add_argument(
        "--poster",
        action="store_true",
        help="also download .jpg poster frames",
    )
    ap.add_argument(
        "--delay",
        type=float,
        default=1.5,
        help="base seconds between requests (default 1.5)",
    )
    ap.add_argument(
        "--jitter",
        type=float,
        default=0.75,
        help="extra random seconds added to delay (default 0.75)",
    )
    ap.add_argument(
        "--backoff",
        type=float,
        default=60.0,
        help="seconds to wait after HTTP 403/429 (default 60)",
    )
    ap.add_argument("--force", action="store_true", help="re-download existing files")
    ap.add_argument(
        "--include-innate",
        action="store_true",
        help="also probe innate abilities (usually 404)",
    )
    args = ap.parse_args()

    index = load_ability_index(include_innate=args.include_innate)

    wanted = [h.strip() for h in args.heroes.split(",") if h.strip()]
    if wanted:
        missing_heroes = [h for h in wanted if h not in index]
        if missing_heroes:
            raise SystemExit(f"unknown heroes (not in hero_abilities.json): {missing_heroes}")
        heroes = wanted
    else:
        heroes = sorted(index.keys())

    exts = []
    if args.fmt in ("webm", "both"):
        exts.append("webm")
    if args.fmt in ("mp4", "both"):
        exts.append("mp4")
    if args.poster:
        exts.append("jpg")

    manifest = load_manifest()
    clips: dict = manifest.setdefault("clips", {})
    missing_set = set(manifest.get("missing") or [])

    ok = skip = miss = fail = 0
    # (hero, ability_key, cdn_name, ext) — cdn_name differs from ability_key
    # for scepter/shard-granted abilities.
    jobs: list[tuple[str, str, str, str]] = []
    for hero in heroes:
        for entry in index[hero]:
            for ext in exts:
                jobs.append((hero, entry["key"], entry["cdn_name"], ext))

    print(
        f"heroes={len(heroes)} jobs={len(jobs)} "
        f"delay={args.delay}+jitter[0,{args.jitter}] "
        f"out={ABILITY_VIDEOS}",
        flush=True,
    )

    for n, (hero, ability, cdn_name, ext) in enumerate(jobs, start=1):
        rel = f"{hero}/{ability}.{ext}"
        dest = ABILITY_VIDEOS / hero / f"{ability}.{ext}"
        url = f"{CDN_BASE}/{hero}/{cdn_name}.{ext}"
        clip_key = f"{hero}/{ability}"

        if (
            not args.force
            and ext != "jpg"
            and clip_key in missing_set
            and not dest.is_file()
        ):
            # Previously probed missing video; skip re-probe unless --force.
            # Still allow poster retry independently.
            if ext in ("webm", "mp4"):
                miss += 1
                print(f"  [{n}/{len(jobs)}] miss(cached) {rel}", flush=True)
                continue

        status = fetch_one(
            url,
            dest,
            delay=args.delay,
            jitter=args.jitter,
            force=args.force,
            backoff=args.backoff,
        )
        if status == "ok":
            ok += 1
            if ext in ("webm", "mp4"):
                cell = clips.setdefault(clip_key, {"hero": hero, "ability": ability})
                cell[ext] = rel.replace("\\", "/")
                cell["bytes_" + ext] = dest.stat().st_size
                missing_set.discard(clip_key)
            print(f"  [{n}/{len(jobs)}] ok {rel} ({dest.stat().st_size} bytes)", flush=True)
        elif status == "skip":
            skip += 1
            if ext in ("webm", "mp4"):
                cell = clips.setdefault(clip_key, {"hero": hero, "ability": ability})
                cell[ext] = rel.replace("\\", "/")
                cell["bytes_" + ext] = dest.stat().st_size
                missing_set.discard(clip_key)
            if n == 1 or n % 25 == 0 or n == len(jobs):
                print(f"  [{n}/{len(jobs)}] skip {rel}", flush=True)
        elif status == "missing":
            miss += 1
            if ext in ("webm", "mp4"):
                missing_set.add(clip_key)
            print(f"  [{n}/{len(jobs)}] miss {rel}", flush=True)
        elif status == "rate_limited":
            fail += 1
            manifest["missing"] = sorted(missing_set)
            save_manifest(manifest)
            print(
                f"stopped early after rate limit: ok={ok} skip={skip} miss={miss} fail={fail}",
                flush=True,
            )
            raise SystemExit(2)
        else:
            fail += 1
            print(f"  [{n}/{len(jobs)}] FAIL {rel}", flush=True)

        if n % 10 == 0 or n == len(jobs):
            manifest["missing"] = sorted(missing_set)
            save_manifest(manifest)

    manifest["missing"] = sorted(missing_set)
    save_manifest(manifest)
    print(
        f"done: downloaded={ok} skipped={skip} missing={miss} failed={fail} "
        f"manifest={MANIFEST}",
        flush=True,
    )


if __name__ == "__main__":
    main()