"""Check Gemini API key, available models, and rate limits."""
import asyncio, os, sys
from dotenv import load_dotenv
import httpx

load_dotenv()
KEY = os.getenv("GOOGLE_API_KEY", "")

BASE = "https://generativelanguage.googleapis.com/v1beta"
MODELS_TO_TEST = [
    "gemini-2.5-flash",
    "gemini-2.5-pro",
    "gemini-2.0-flash",
    "gemini-1.5-flash",
    "gemini-1.5-pro",
]


async def main():
    if not KEY:
        print("[X] GOOGLE_API_KEY not set in .env")
        sys.exit(1)
    print(f"Key: ...{KEY[-6:]}\n")

    async with httpx.AsyncClient(timeout=20) as c:

        # 1) List all available models
        print("=" * 55)
        print("AVAILABLE MODELS")
        print("=" * 55)
        r = await c.get(f"{BASE}/models?key={KEY}")
        if r.status_code != 200:
            print(f"[X] Could not list models: {r.status_code} {r.text[:200]}")
        else:
            models = r.json().get("models", [])
            gen_models = [m for m in models if "generateContent" in m.get("supportedGenerationMethods", [])]
            for m in gen_models:
                name = m["name"].replace("models/", "")
                limit = m.get("description", "")[:60]
                print(f"  {name:<35} rpm={m.get('rpmLimit','?'):>6}  tpm={m.get('tpmLimit','?'):>10}")

        # 2) Test each target model with a tiny call
        print("\n" + "=" * 55)
        print("MODEL PING TEST")
        print("=" * 55)
        body = {"contents": [{"parts": [{"text": "Say OK"}]}],
                "generationConfig": {"maxOutputTokens": 5}}
        for model in MODELS_TO_TEST:
            url = f"{BASE}/models/{model}:generateContent?key={KEY}"
            r = await c.post(url, json=body)
            if r.status_code == 200:
                status = "[OK]"
            elif r.status_code == 429:
                status = "[429 rate limit]"
            elif r.status_code == 404:
                status = "[404 not found]"
            elif r.status_code == 403:
                status = "[403 denied]"
            else:
                status = f"[{r.status_code}]"
            print(f"  {model:<30} {status}")

        # 3) Quota info from a real call on the working model
        print("\n" + "=" * 55)
        print("RATE LIMITS (from API metadata)")
        print("=" * 55)
        r = await c.get(f"{BASE}/models/gemini-2.5-flash?key={KEY}")
        if r.status_code == 200:
            m = r.json()
            print(f"  Model        : {m.get('displayName','')}")
            print(f"  Input limit  : {m.get('inputTokenLimit','?'):,} tokens")
            print(f"  Output limit : {m.get('outputTokenLimit','?'):,} tokens")
            print(f"  RPM limit    : {m.get('rpmLimit', 'see https://ai.dev/rate-limit')}")
            print(f"  TPM limit    : {m.get('tpmLimit', 'see https://ai.dev/rate-limit')}")
            print(f"\n  Free tier caps: 15 RPM / 1,000,000 TPM / 1,500 RPD")
            print(f"  Full limits  : https://ai.google.dev/gemini-api/docs/rate-limits")


asyncio.run(main())