"""Cache the architecture of licence-gated models the demo cannot read.

Some of the best-known LLMs -- Llama and Gemma -- are licence-gated on the
Hugging Face Hub, so an anonymous demo cannot read their `config.json` and has
to refuse them. Serving them with a maintainer's own token would work, but it
would mean using one person's licence acceptance on behalf of every visitor.

Instead this script reads the architecture from a public, unquantised mirror of
the same weights and caches the handful of integers WattGPU actually needs.
Those numbers -- layer count, hidden size, head counts, parameter count -- are
published in the model cards and papers, so caching them redistributes nothing.

Each entry records the mirror it came from, so any value can be traced and
re-checked. Run this again to refresh:

    python scripts/cache_gated_models.py
"""

from __future__ import annotations

import argparse
import json
import os
import sys

REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, REPO_ROOT)

from wattgpu_demo.hf_models import (  # noqa: E402
    HF_ENDPOINT,
    _get_json,
    detect_quantisation,
    llm_info_from_config,
)

# Gated model -> a public repository holding the same architecture at full
# precision. Mirrors are checked for quantisation before anything is cached.
MIRRORS = {
    "meta-llama/Llama-3.2-1B": "unsloth/Llama-3.2-1B",
    "meta-llama/Llama-3.2-1B-Instruct": "unsloth/Llama-3.2-1B-Instruct",
    "meta-llama/Meta-Llama-3-8B-Instruct": "NousResearch/Meta-Llama-3-8B-Instruct",
    "google/gemma-3-1b-it": "unsloth/gemma-3-1b-it",
    "google/gemma-3-270m": "unsloth/gemma-3-270m",
}

DEFAULT_OUT = os.path.join(REPO_ROOT, "data", "gated_llms.json")


def fetch_entry(canonical: str, mirror: str) -> dict:
    config = _get_json(f"{HF_ENDPOINT}/{mirror}/raw/main/config.json", mirror)

    quantisation = detect_quantisation(mirror, config)
    if quantisation:
        raise ValueError(
            f"{mirror} is quantised ({quantisation}); it cannot stand in for {canonical}")

    info = _get_json(f"{HF_ENDPOINT}/api/models/{mirror}", mirror)
    total = ((info or {}).get("safetensors") or {}).get("total")
    if not isinstance(total, (int, float)) or total <= 0:
        raise ValueError(f"{mirror} publishes no safetensors index")

    # Reuse the same adapter the live path uses, so a cached model and a
    # fetched one go through identical validation.
    llm = llm_info_from_config(canonical, config, float(total) / 1e9, "published config")
    return {
        "model_type": llm.model_type,
        "num_layers": llm.num_layers,
        "hidden_size": llm.hidden_size,
        "num_attention_heads": llm.num_attention_heads,
        "num_key_value_heads": llm.num_key_value_heads,
        "total_b_params": round(llm.total_b_params, 6),
        "architectures": llm.architectures,
        "max_position_embeddings": llm.max_position_embeddings,
        "torch_dtype": llm.torch_dtype,
        "source_mirror": mirror,
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out", default=DEFAULT_OUT)
    args = parser.parse_args()

    entries: dict[str, dict] = {}
    for canonical, mirror in MIRRORS.items():
        try:
            entries[canonical] = fetch_entry(canonical, mirror)
            e = entries[canonical]
            print(f"  {canonical}: {e['total_b_params']:.3f} B, "
                  f"{e['num_layers']} layers, via {mirror}")
        except Exception as exc:  # noqa: BLE001 - report and keep going
            print(f"  SKIPPED {canonical}: {exc}")

    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    with open(args.out, "w") as fh:
        json.dump(entries, fh, indent=2, sort_keys=True)
    print(f"\nwrote {len(entries)} gated-model architectures to {args.out}")
    return 0 if entries else 1


if __name__ == "__main__":
    sys.exit(main())