#!/usr/bin/env python3
"""Re-runnable xAI model discovery.

GET /v1/models                      -> sanitized sources/xai/models-api-raw.json (mixed catalogue, 12 ids on 2026-09-19)
GET /v1/language-models             -> sources/xai/language-models-raw.json
GET /v1/image-generation-models     -> sources/xai/image-generation-models-raw.json
GET /v1/video-generation-models     -> sources/xai/video-generation-models-raw.json
GET /v1/embedding-models            -> sources/xai/embedding-models-raw.json
Diff of /v1/models against the previous file -> sources/xai/models-diff.json, prints NEW / REMOVED / CHANGED ids
(alias changes and price changes count as CHANGED; prices are ticks, 1 USD = 1e10).

Usage:  python3 scripts/discover_xai_models.py [--dry-run] [--base https://us.api.x.ai]
Auth comes from .env via scripts/live.py (never printed). Each call is logged to reports/live-requests.jsonl.
All these endpoints are free (no usage object is returned).
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import live  # noqa: E402

ROOT = Path(__file__).resolve().parents[1]
SRC = ROOT / "sources" / "xai"
RAW = SRC / "models-api-raw.json"
DIFF = SRC / "models-diff.json"
TYPED = {
    "/v1/language-models": SRC / "language-models-raw.json",
    "/v1/image-generation-models": SRC / "image-generation-models-raw.json",
    "/v1/video-generation-models": SRC / "video-generation-models-raw.json",
    "/v1/embedding-models": SRC / "embedding-models-raw.json",
}


def get(path: str, note: str) -> dict:
    st, body, _ = live.xai_request("GET", path, note=note)
    if st != 200:
        raise SystemExit(f"GET {path} -> HTTP {st}: {json.dumps(body)[:300]}")
    return body


def index(doc: dict) -> dict[str, dict]:
    return {m["id"]: m for m in doc.get("data", [])}


def fmt_price(m: dict) -> str:
    if m.get("prompt_text_token_price") is not None:
        return (f"${m['prompt_text_token_price'] / 1e4:.2f}/{m['cached_prompt_text_token_price'] / 1e4:.2f}/"
                f"{m['completion_text_token_price'] / 1e4:.2f} per 1M, ctx={m.get('context_length')}")
    if m.get("image_price") is not None:
        return f"${m['image_price'] / 1e10:.3f}/image"
    return "(no price fields)"


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--dry-run", action="store_true", help="fetch and diff, do not write")
    ap.add_argument("--base", default=None, help="override base URL (e.g. https://us.api.x.ai); typed lists + write skipped")
    args = ap.parse_args()
    if args.base:
        os.environ["XAI_BASE_URL"] = args.base
        live.XAI_BASE = args.base

    prev = json.loads(RAW.read_text()) if RAW.exists() and not args.base else {"data": []}
    cur = get("/v1/models", f"discover models{' @' + args.base if args.base else ''}")
    cur["_retrieved_at"] = live.now()
    pi, ci = index(prev), index(cur)
    new = sorted(set(ci) - set(pi))
    removed = sorted(set(pi) - set(ci))
    changed = sorted(i for i in set(ci) & set(pi) if ci[i] != pi[i])
    print(f"models: {len(ci)} (previous {len(pi)}) base={args.base or live.XAI_BASE}")
    for i in new:
        print(f"NEW      {i}  {fmt_price(ci[i])}  aliases={ci[i].get('aliases')}")
    for i in removed:
        print(f"REMOVED  {i}")
    for i in changed:
        keys = sorted(k for k in set(ci[i]) | set(pi[i]) if ci[i].get(k) != pi[i].get(k))
        print(f"CHANGED  {i}  fields={keys}")
    if not (new or removed or changed):
        print("no changes")
    if args.base:
        for m in cur["data"]:
            print(f"  {m['id']}: {fmt_price(m)}")
        return

    typed: dict[str, dict] = {}
    for path, dest in TYPED.items():
        typed[path] = get(path, "discover typed model list")
        print(f"{path}: {len(typed[path].get('models', []))} models")

    diff = {"retrieved_at": cur["_retrieved_at"], "previous_count": len(pi), "current_count": len(ci), "new": new, "removed": removed,
            "changed": [{"id": i, "before": pi[i], "after": ci[i]} for i in changed],
            "typed_counts": {p: len(d.get("models", [])) for p, d in typed.items()}}
    if args.dry_run:
        return
    live.save_sanitized({k: v for k, v in cur.items() if not k.startswith("_")}, RAW)
    for path, dest in TYPED.items():
        live.save_sanitized(typed[path], dest)
    live.save_sanitized(diff, DIFF)
    print(f"wrote {RAW.relative_to(ROOT)}, typed raws and {DIFF.relative_to(ROOT)}")


if __name__ == "__main__":
    main()