#!/usr/bin/env python3
"""Refresh bundled Endgame team snapshots from the credited Prydwen dataset."""

from __future__ import annotations

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(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)


def normalized(value: str) -> str:
    return "".join(char.lower() for char in value if char.isalnum())


def latest_dataset_version(config: dict[str, dict]) -> tuple[str, dict]:
    available = []
    for version, metadata in config.items():
        date_text = metadata.get("collect_date")
        if not version[0].isdigit() or not date_text:
            continue
        try:
            collected = dt.datetime.strptime(date_text, "%d/%m/%Y").date()
        except ValueError:
            continue
        if all(isinstance(metadata.get(mode), dict) for mode in MODES):
            available.append((collected, version, metadata))
    if not available:
        raise RuntimeError("No dated source snapshot has all four Endgame modes")
    _, version, metadata = max(available)
    return version, metadata


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)
    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 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}")
    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("--mode", choices=("all", *MODES), default="all", help="refresh all modes or one Endgame mode")
    args = parser.parse_args()

    config = fetch_json(f"{DATASET_ROOT}/config.json")
    if not isinstance(config, dict):
        raise RuntimeError("Unexpected source config shape")
    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:
            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",
            "url": "https://huggingface.co/datasets/LvlUrArti/MocDataProcessed",
            "license": "MIT",
            "attribution": "Processed data shown on Prydwen; credit LvlUrArti",
            "dataset_version": version,
            "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,
        "modes": modes,
    }
    RESOURCE_PATH.write_text(json.dumps(output, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(f"Updated {', '.join(selected_modes)} from dataset {version} ({len(modes)} modes)")


if __name__ == "__main__":
    main()