"""Test which Gemini/Gemma models are available for your API key."""

import google.generativeai as genai
import os
from dotenv import load_dotenv

load_dotenv()

GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
if not GEMINI_API_KEY:
    print("ERROR: GEMINI_API_KEY not found in .env")
    exit(1)

genai.configure(api_key=GEMINI_API_KEY)

# Models to test
MODELS = [
    "gemini-2.0-flash",
    "gemini-2.0-flash-lite",
    "gemma-3-1b-it",
    "gemma-3-4b-it",
    "gemma-3-12b-it",
    "gemma-3-27b-it",
]

TEST_PROMPT = "Say 'Hello' in one word."

print("=" * 60)
print("Testing Gemini/Gemma Models Availability")
print("=" * 60)

available = []
unavailable = []

for model_name in MODELS:
    try:
        model = genai.GenerativeModel(model_name)
        response = model.generate_content(TEST_PROMPT)
        text = response.text.strip()[:50]
        print(f"[OK]  {model_name:25} -> {text}")
        available.append(model_name)
    except Exception as e:
        error_msg = str(e)[:60]
        print(f"[ERR] {model_name:25} -> {error_msg}")
        unavailable.append(model_name)

print("\n" + "=" * 60)
print("SUMMARY")
print("=" * 60)
print(f"\nAvailable ({len(available)}):")
for m in available:
    print(f"  - {m}")

print(f"\nUnavailable ({len(unavailable)}):")
for m in unavailable:
    print(f"  - {m}")

# Calculate combined RPM
rpm_map = {
    "gemini-2.0-flash": 15,
    "gemini-2.0-flash-lite": 30,
    "gemma-3-1b-it": 30,
    "gemma-3-4b-it": 30,
    "gemma-3-12b-it": 30,
    "gemma-3-27b-it": 30,
}

total_rpm = sum(rpm_map.get(m, 0) for m in available)
print(f"\nCombined RPM capacity: {total_rpm} RPM")
print(f"With 50% cache hit rate: ~{total_rpm * 2} effective RPM")