from google import genai
import os
from dotenv import load_dotenv

load_dotenv()

api_key = os.getenv("GEMINI_API_KEY")

if not api_key:
    print("API Key not found")
else:
    client = genai.Client(api_key=api_key)
    try:
        print("Testing models...", flush=True)
        pager = client.models.list()
        
        with open("working_models.txt", "w") as f:
            for model in pager:
                print(f"Testing {model.name}...", flush=True)
                # We want stable Gemini models
                if "gemini" in model.name and "vision" not in model.name:
                     try:
                        response = client.models.generate_content(
                            model=model.name,
                            contents="Hello"
                        )
                        print(f"✅ SUCCESS: {model.name}", flush=True)
                        f.write(f"{model.name}\n")
                     except Exception as e:
                        print(f"❌ FAILED: {model.name} - {e}", flush=True)
                else:
                     print(f"Skipping {model.name}", flush=True)
            
    except Exception as e:
        print(f"Error listing models: {e}", flush=True)