#!/usr/bin/env python3
"""Regenerate src/providers_data.rs from Mastra's provider-registry.json.

Fetches the canonical provider registry (base URLs, env vars, model lists)
from the mastra-ai/mastra GitHub repo and emits a Rust static table.
Curated overrides fix poor auto-picked defaults and fill missing base URLs.

Usage: python3 scripts/gen_providers.py
"""

import json
import re
import urllib.request

REGISTRY_URL = (
    "https://raw.githubusercontent.com/mastra-ai/mastra/main/"
    "packages/core/src/llm/model/provider-registry.json"
)
OUT = "src/llm/providers_data.rs"

BAD = re.compile(
    r"(embedding|realtime|^.*image|rerank|^tts|speech|whisper|audio|transcri|"
    r"translate|^text-|sound|multimodal|^vision|guard|safeguard|ocr)",
    re.I,
)

CURATED = {
    "openai": (None, "gpt-4o-mini"),
    "anthropic": (None, "claude-sonnet-4-20250514"),
    "google": (None, "gemini-2.5-flash"),
    "groq": (None, "llama-3.3-70b-versatile"),
    "xai": ("https://api.x.ai/v1", "grok-4.3"),
    "cerebras": ("https://api.cerebras.ai/v1", "gpt-oss-120b"),
    "deepinfra": ("https://api.deepinfra.com/v1/openai", "Qwen/Qwen3-32B"),
    "perplexity": ("https://api.perplexity.ai", "sonar-pro"),
    "togetherai": ("https://api.together.xyz/v1", "MiniMaxAI/MiniMax-M2.5"),
    "moonshotai": (None, "kimi-k2-turbo-preview"),
    "moonshotai-cn": (None, "kimi-k2-turbo-preview"),
    "zai": (None, "glm-4.6"),
    "alibaba": (None, "qwen-plus"),
    "sakana": (None, "fugu-ultra"),
    "upstage": (None, "solar-pro2"),
    "deepseek": (None, "deepseek-chat"),
    "opencode": (None, "big-pickle"),
    "opencode-go": (None, "deepseek-v4-flash"),
    "vercel": ("https://ai.gateway.vercel.ai/v1", None),
}


def fetch() -> dict:
    with urllib.request.urlopen(REGISTRY_URL) as r:
        return json.load(r)


def pick_model(v: dict) -> str:
    models = v.get("models", [])
    depr = set(v.get("deprecatedModels", []))
    for m in models:
        if m in depr:
            continue
        if BAD.search(m):
            continue
        return m
    return models[0] if models else ""


def pick_models(v: dict, cap: int = 300) -> list:
    """All non-deprecated chat-ish model names, sorted, capped for size."""
    depr = set(v.get("deprecatedModels", []))
    out = []
    for m in v.get("models", []):
        if m in depr or BAD.search(m):
            continue
        if m not in out:
            out.append(m)
    out.sort()
    return out[:cap]


def fmt_for(npm) -> str:
    if not npm:
        return "Compatible"
    if "anthropic" in npm:
        return "Anthropic"
    if "google" in npm and "vertex" not in npm:
        return "Gemini"
    if "bedrock" in npm or "vertex" in npm:
        return "Bedrock"
    return "Compatible"


def esc(s: str) -> str:
    return s.replace("\\", "\\\\").replace('"', '\\"')


def main() -> None:
    data = fetch()
    provs = data["providers"]
    rows = []
    model_lists = {}
    for key, v in sorted(provs.items()):
        env = v.get("apiKeyEnvVar") or ""
        if isinstance(env, list):
            env = ",".join(env)
        base = v.get("url") or ""
        model = pick_model(v)
        curated = CURATED.get(key)
        if curated:
            if curated[0]:
                base = curated[0]
            if curated[1]:
                model = curated[1]
        rows.append((key, env, fmt_for(v.get("npm")), base, model))
        model_lists[key] = pick_models(v)

    lines = [
        "// Generated by scripts/gen_providers.py from Mastra's provider-registry.json",
        "// (https://github.com/mastra-ai/mastra). Do not edit by hand.",
        "",
        "use super::ApiFormat;",
        "",
        "pub struct ProviderRow {",
        "    pub name: &'static str,",
        "    pub format: ApiFormat,",
        "    pub env: &'static str,",
        "    pub base_url: &'static str,",
        "    pub default_model: &'static str,",
        "    pub requires_key: bool,",
        "}",
        "",
        "pub const PROVIDER_ROWS: &[ProviderRow] = &[",
    ]
    fmt_rs = {
        "Compatible": "ApiFormat::Compatible",
        "Anthropic": "ApiFormat::Anthropic",
        "Gemini": "ApiFormat::Gemini",
        "Bedrock": "ApiFormat::Bedrock",
    }
    for name, env, fmt, base, model in rows:
        rk = "true" if env else "false"
        lines.append(
            f'    ProviderRow {{ name: "{esc(name)}", format: {fmt_rs[fmt]}, '
            f'env: "{esc(env)}", base_url: "{esc(base)}", '
            f'default_model: "{esc(model)}", requires_key: {rk} }},'
        )
    lines.append("];")
    lines.append("")
    lines.append("pub const MODEL_LISTS: &[(&str, &[&str])] = &[")
    for key in sorted(model_lists):
        models = model_lists[key]
        if not models:
            continue
        joined = ", ".join(f'"{esc(m)}"' for m in models)
        lines.append(f'    ("{esc(key)}", &[{joined}]),')
    lines.append("];")
    lines.append("")
    with open(OUT, "w") as f:
        f.write("\n".join(lines) + "\n")
    print(f"wrote {OUT}: {len(rows)} providers")


if __name__ == "__main__":
    main()