Update GaCy-GO

tools/update_prydwen_endgame_teams.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Refresh the MoC recommendation snapshot from the credited Prydwen dataset."""
+"""Refresh bundled Endgame team snapshots from the credited Prydwen dataset."""
 
 from __future__ import annotations
 
@@ -7,19 +7,21 @@ import argparse
 import datetime as dt
 import json
 import pathlib
+import re
 import urllib.request
 
 
 DATASET_ROOT = "https://huggingface.co/datasets/LvlUrArti/MocDataProcessed/resolve/main"
+DATASET_API_ROOT = "https://huggingface.co/api/datasets/LvlUrArti/MocDataProcessed"
 REPOSITORY_ROOT = pathlib.Path(__file__).resolve().parents[1]
 RESOURCE_PATH = REPOSITORY_ROOT / "resources" / "PrydwenEndgameTeams.json"
 AVATAR_CONFIG_PATH = REPOSITORY_ROOT / "resources" / "AvatarConfig.json"
+MODES = ("moc", "pf", "as", "aa")
+DEFAULT_REQUEST_TYPE_ROUTES = {"2": "moc", "3": "pf", "4": "as", "5": "aa"}
 
 
-def fetch_json(path: str) -> object:
-    request = urllib.request.Request(
-        f"{DATASET_ROOT}/{path}", headers={"User-Agent": "GaC-hkrpg-go-recommendations/1.0"}
-    )
+def fetch_json(url: str) -> object:
+    request = urllib.request.Request(url, headers={"User-Agent": "GaC-hkrpg-go-recommendations/1.0"})
     with urllib.request.urlopen(request, timeout=30) as response:
         return json.load(response)
 
@@ -31,72 +33,105 @@ def normalized(value: str) -> str:
 def latest_dataset_version(config: dict[str, dict]) -> tuple[str, dict]:
     available = []
     for version, metadata in config.items():
-        mode = metadata.get("moc")
         date_text = metadata.get("collect_date")
-        if not version[0].isdigit() or not isinstance(mode, dict) or not date_text:
+        if not version[0].isdigit() or not date_text:
             continue
         try:
             collected = dt.datetime.strptime(date_text, "%d/%m/%Y").date()
         except ValueError:
             continue
-        available.append((collected, version, mode))
+        if all(isinstance(metadata.get(mode), dict) for mode in MODES):
+            available.append((collected, version, metadata))
     if not available:
-        raise RuntimeError("No dated MoC snapshot found in the source dataset config")
-    _, version, mode = max(available)
-    return version, mode
+        raise RuntimeError("No dated source snapshot has all four Endgame modes")
+    _, version, metadata = max(available)
+    return version, metadata
 
 
-def avatar_tags() -> set[str]:
+def avatar_index(aliases: dict[str, str]) -> tuple[set[str], dict[str, str]]:
     with AVATAR_CONFIG_PATH.open(encoding="utf-8") as source:
         config = json.load(source)
-    return {
+    tags = {
         normalized(entry["AvatarVOTag"])
         for entry in config.get("avatar_config", [])
         if entry.get("AvatarVOTag")
     }
+    for slug, client_tag in aliases.items():
+        if normalized(client_tag) in tags:
+            tags.add(normalized(slug))
+    return tags, aliases
 
 
-def stage_team(version: str, floor: int, stage: int, tags: set[str], aliases: dict[str, str]) -> dict:
-    path = f"{version}/moc/comps/{floor}-{stage}_combined.json"
-    rows = fetch_json(path)
+def dataset_stage_files(version: str, mode: str, tree: list[dict]) -> list[tuple[str, str]]:
+    if not isinstance(tree, list):
+        raise RuntimeError("Unexpected Hugging Face dataset tree shape")
+    prefix = f"{version}/{mode}/comps/"
+    stages = []
+    for entry in tree:
+        path = entry.get("path", "")
+        if entry.get("type") != "file" or not path.startswith(prefix):
+            continue
+        match = re.fullmatch(r"(\d+-\d+)_combined\.json", pathlib.PurePosixPath(path).name)
+        if match:
+            stages.append((match.group(1), path))
+    if not stages:
+        raise RuntimeError(f"No stage team files found for {version}/{mode}")
+    return sorted(stages)
+
+
+def best_compatible_team(path: str, tags: set[str], aliases: dict[str, str]) -> dict:
+    rows = fetch_json(f"{DATASET_ROOT}/{path}")
     if not isinstance(rows, list):
         raise RuntimeError(f"Unexpected dataset shape in {path}")
-    top = next((row for row in rows if row.get("rank") == 1), None)
-    if top is None:
-        raise RuntimeError(f"No rank-1 lineup found in {path}")
-    characters = [top.get(f"char_{slot}") for slot in ("one", "two", "three", "four")]
-    if any(not isinstance(slug, str) or not slug for slug in characters):
-        raise RuntimeError(f"Incomplete rank-1 lineup in {path}")
-    unresolved = [slug for slug in characters if normalized(aliases.get(slug, slug)) not in tags]
-    if unresolved:
-        raise RuntimeError(
-            f"Source slugs do not map to this Client's AvatarConfig: {', '.join(unresolved)}; "
-            "add their source-slug to AvatarVOTag aliases before refreshing"
-        )
-    return {"stage": stage, "rank": 1, "characters": characters}
+    for row in sorted(rows, key=lambda value: value.get("rank", 1 << 30)):
+        characters = [row.get(f"char_{slot}") for slot in ("one", "two", "three", "four")]
+        if any(not isinstance(slug, str) or not slug for slug in characters):
+            continue
+        if len({normalized(slug) for slug in characters}) != 4:
+            continue
+        if any(normalized(aliases.get(slug, slug)) not in tags for slug in characters):
+            continue
+        return {"rank": int(row.get("rank", 0)), "characters": characters}
+    raise RuntimeError(f"No four-character lineup in {path} maps to this Client's AvatarConfig")
 
 
 def main() -> None:
     parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--floor", type=int, required=True, help="Prydwen MoC floor shown by the Client")
-    parser.add_argument("--stage-count", type=int, required=True, help="number of stage teams to fetch")
+    parser.add_argument("--mode", choices=("all", *MODES), default="all", help="refresh all modes or one Endgame mode")
     args = parser.parse_args()
-    if args.floor <= 0 or args.stage_count <= 0:
-        parser.error("--floor and --stage-count must be positive")
 
-    config = fetch_json("config.json")
+    config = fetch_json(f"{DATASET_ROOT}/config.json")
     if not isinstance(config, dict):
         raise RuntimeError("Unexpected source config shape")
-    version, cycle = latest_dataset_version(config)
-    aliases = {}
+    version, version_data = latest_dataset_version(config)
+    tree = fetch_json(f"{DATASET_API_ROOT}/tree/main?recursive=true&expand=false")
+    if not isinstance(tree, list):
+        raise RuntimeError("Unexpected Hugging Face dataset tree shape")
+
+    previous = {}
     if RESOURCE_PATH.exists():
         with RESOURCE_PATH.open(encoding="utf-8") as source:
-            aliases = json.load(source).get("avatar_aliases", {})
-    tags = avatar_tags()
-    stages = [
-        stage_team(version, args.floor, stage, tags, aliases)
-        for stage in range(1, args.stage_count + 1)
-    ]
+            previous = json.load(source)
+    aliases = previous.get("avatar_aliases", {})
+    tags, aliases = avatar_index(aliases)
+    modes = previous.get("modes", {})
+    selected_modes = MODES if args.mode == "all" else (args.mode,)
+    for mode in selected_modes:
+        stage_files = dataset_stage_files(version, mode, tree)
+        stages = []
+        for stage, path in stage_files:
+            team = best_compatible_team(path, tags, aliases)
+            stages.append({"stage": stage, **team})
+        metadata = version_data[mode]
+        modes[mode] = {
+            "dataset_version": version,
+            "challenge_version": metadata.get("ver"),
+            "challenge_name": metadata.get("name"),
+            "floor": max(int(stage["stage"].split("-")[0]) for stage in stages),
+            "stages": stages,
+        }
+
+    routes = previous.get("protocol", {}).get("request_type_routes", DEFAULT_REQUEST_TYPE_ROUTES)
     output = {
         "source": {
             "name": "LvlUrArti/MocDataProcessed",
@@ -104,17 +139,17 @@ def main() -> None:
             "license": "MIT",
             "attribution": "Processed data shown on Prydwen; credit LvlUrArti",
             "dataset_version": version,
-            "collect_date": dt.datetime.strptime(config[version]["collect_date"], "%d/%m/%Y").date().isoformat(),
-            "challenge_version": cycle.get("ver"),
-            "challenge_name": cycle.get("name"),
+            "collect_date": dt.datetime.strptime(version_data["collect_date"], "%d/%m/%Y").date().isoformat(),
+        },
+        "protocol": {
+            "request_type_routes": routes,
+            "mapping_note": "Type 2 is retained for MoC; Types 3-5 are inferred mode routes pending Client packet confirmation.",
         },
         "avatar_aliases": aliases,
-        "mode": "moc",
-        "floor": args.floor,
-        "stages": stages,
+        "modes": modes,
     }
     RESOURCE_PATH.write_text(json.dumps(output, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
-    print(f"Updated {RESOURCE_PATH.relative_to(REPOSITORY_ROOT)} from dataset {version}: {len(stages)} stages")
+    print(f"Updated {', '.join(selected_modes)} from dataset {version} ({len(modes)} modes)")
 
 
 if __name__ == "__main__":