#!/usr/bin/env python3
"""Re-runnable Gemini model discovery (Gemini Developer API).

GET /v1beta/models (paginated: pageSize + pageToken) -> sanitized sources/gemini/models-api-raw.json
GET /v1/models      (stable surface)                 -> sources/gemini/models-v1-raw.json
Diff against the previous v1beta file -> sources/gemini/models-diff.json, prints NEW / REMOVED / CHANGED
(limits, supportedGenerationMethods, thinking, version, defaults) and the v1-vs-v1beta id difference.

Usage:  python3 scripts/discover_gemini_models.py [--page-size N] [--dry-run] [--skip-v1]
Auth comes from .env via scripts/live.py (x-goog-api-key header, never in the URL, never printed).
Each call is logged to reports/live-requests.jsonl (path is normalised to /v1beta/models).
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

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

ROOT = Path(__file__).resolve().parents[1]
RAW = ROOT / "sources" / "gemini" / "models-api-raw.json"
RAW_V1 = ROOT / "sources" / "gemini" / "models-v1-raw.json"
DIFF = ROOT / "sources" / "gemini" / "models-diff.json"

# Fields whose change is worth flagging as CHANGED (order irrelevant for lists).
WATCH = ("inputTokenLimit", "outputTokenLimit", "supportedGenerationMethods", "thinking", "version",
         "temperature", "topP", "topK", "maxTemperature", "displayName", "description")


def fetch_all(version: str, page_size: int) -> dict:
    models: list[dict] = []
    token: str | None = None
    pages = 0
    while True:
        q = f"?pageSize={page_size}" + (f"&pageToken={token}" if token else "")
        st, body, _ = live.gemini_request("GET", f"/{version}/models{q}", note=f"discover models {version} page {pages + 1}")
        if st != 200:
            raise SystemExit(f"GET /{version}/models -> HTTP {st}: {json.dumps(body)[:300]}")
        models.extend(body.get("models", []))
        pages += 1
        token = body.get("nextPageToken")
        if not token:
            break
    return {"models": models, "_pages": pages, "_retrieved_at": live.now(), "_version": version}


def index(doc: dict) -> dict[str, dict]:
    return {m["name"].removeprefix("models/"): m for m in doc.get("models", [])}


def changed_fields(a: dict, b: dict) -> list[str]:
    out = []
    for k in WATCH:
        va, vb = a.get(k), b.get(k)
        if isinstance(va, list) and isinstance(vb, list):
            if sorted(va) != sorted(vb):
                out.append(k)
        elif va != vb:
            out.append(k)
    return out


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--page-size", type=int, default=50, help="pageSize query param (docs: default 50, max 1000)")
    ap.add_argument("--dry-run", action="store_true", help="fetch and diff, do not write")
    ap.add_argument("--skip-v1", action="store_true", help="do not call GET /v1/models")
    args = ap.parse_args()

    prev = json.loads(RAW.read_text()) if RAW.exists() else {"models": []}
    cur = fetch_all("v1beta", args.page_size)
    pi, ci = index(prev), index(cur)
    new = sorted(set(ci) - set(pi))
    removed = sorted(set(pi) - set(ci))
    changed = [{"id": i, "fields": changed_fields(pi[i], ci[i]), "before": {k: pi[i].get(k) for k in WATCH}, "after": {k: ci[i].get(k) for k in WATCH}}
               for i in sorted(set(ci) & set(pi)) if changed_fields(pi[i], ci[i])]

    v1_ids: list[str] = []
    only_v1beta: list[str] = []
    only_v1: list[str] = []
    cur_v1 = None
    if not args.skip_v1:
        cur_v1 = fetch_all("v1", args.page_size)
        vi = index(cur_v1)
        v1_ids = sorted(vi)
        only_v1beta = sorted(set(ci) - set(vi))
        only_v1 = sorted(set(vi) - set(ci))

    diff = {"retrieved_at": cur["_retrieved_at"], "previous_count": len(pi), "current_count": len(ci),
            "pages_fetched_v1beta": cur["_pages"], "page_size": args.page_size,
            "new": new, "removed": removed, "changed": changed,
            "v1": {"count": len(v1_ids), "only_in_v1beta": only_v1beta, "only_in_v1": only_v1} if cur_v1 else None}

    print(f"v1beta models: {len(ci)} (previous {len(pi)}), pages={cur['_pages']} (pageSize={args.page_size})")
    for i in new:
        print(f"NEW      {i}  in={ci[i].get('inputTokenLimit')} out={ci[i].get('outputTokenLimit')} methods={ci[i].get('supportedGenerationMethods')}")
    for i in removed:
        print(f"REMOVED  {i}")
    for c in changed:
        print(f"CHANGED  {c['id']}  fields={c['fields']}")
    if not (new or removed or changed):
        print("no changes vs previous v1beta snapshot")
    if cur_v1:
        print(f"v1 models: {len(v1_ids)}; only in v1beta: {only_v1beta or 'none'}; only in v1: {only_v1 or 'none'}")
    if args.dry_run:
        return
    live.save_sanitized({"models": cur["models"]}, RAW)
    if cur_v1:
        live.save_sanitized({"models": cur_v1["models"]}, RAW_V1)
    live.save_sanitized(diff, DIFF)
    print(f"wrote {RAW.relative_to(ROOT)}" + (f", {RAW_V1.relative_to(ROOT)}" if cur_v1 else "") + f" and {DIFF.relative_to(ROOT)}")


if __name__ == "__main__":
    main()