#!/usr/bin/env python3
"""
Test API access to Google Gemini APIs (NanoBanana and Veo3).

Usage:
    python test_api_access.py YOUR_API_KEY

    or set environment variable:
    export GEMINI_API_KEY="your-key"
    python test_api_access.py
"""

import sys
import os
import requests
import json

BASE_URL = "https://generativelanguage.googleapis.com/v1beta"

# API endpoints
ENDPOINTS = {
    "NanoBanana (Image)": {
        "url": f"{BASE_URL}/models/gemini-2.5-flash-image:generateContent",
        "payload": {
            "contents": [{"parts": [{"text": "A red dot"}]}],
            "generationConfig": {"responseModalities": ["TEXT", "IMAGE"]}
        }
    },
    "Veo3 (Video)": {
        "url": f"{BASE_URL}/models/veo-3.1-generate-preview:predictLongRunning",
        "payload": {
            "instances": [{"prompt": "A red dot"}],
            "parameters": {"aspectRatio": "16:9"}
        }
    }
}


def test_api(name: str, config: dict, api_key: str) -> tuple[bool, str]:
    """Test a single API endpoint."""
    headers = {
        "x-goog-api-key": api_key,
        "Content-Type": "application/json"
    }

    try:
        response = requests.post(
            config["url"],
            headers=headers,
            json=config["payload"],
            timeout=30
        )

        data = response.json()

        if response.status_code == 200:
            return True, "Access granted"
        elif "error" in data:
            error = data["error"]
            code = error.get("code", "?")
            message = error.get("message", "Unknown error")[:100]
            return False, f"Error {code}: {message}"
        else:
            return False, f"HTTP {response.status_code}"

    except requests.exceptions.Timeout:
        return False, "Request timeout"
    except requests.exceptions.RequestException as e:
        return False, f"Network error: {str(e)[:50]}"
    except json.JSONDecodeError:
        return False, "Invalid JSON response"


def main():
    # Get API key
    api_key = None
    if len(sys.argv) > 1:
        api_key = sys.argv[1]
    else:
        api_key = os.environ.get("GEMINI_API_KEY")

    if not api_key:
        print("Usage: python test_api_access.py YOUR_API_KEY")
        print("   or: export GEMINI_API_KEY='your-key' && python test_api_access.py")
        print("\nGet your API key at: https://aistudio.google.com/apikey")
        sys.exit(1)

    print("=" * 50)
    print("Testing Google Gemini API Access")
    print("=" * 50)
    print(f"API Key: {api_key[:8]}...{api_key[-4:]}")
    print("-" * 50)

    results = {}
    for name, config in ENDPOINTS.items():
        print(f"\nTesting {name}...")
        success, message = test_api(name, config, api_key)
        results[name] = (success, message)

        status = "✓ OK" if success else "✗ FAIL"
        print(f"  {status}: {message}")

    print("\n" + "=" * 50)
    print("Summary:")
    print("-" * 50)

    all_ok = True
    for name, (success, message) in results.items():
        status = "✓" if success else "✗"
        print(f"  {status} {name}")
        if not success:
            all_ok = False

    print("=" * 50)

    if all_ok:
        print("\nAll APIs accessible! You're ready to go.")
    else:
        print("\nSome APIs are not accessible.")
        print("Check your API key and permissions at:")
        print("https://aistudio.google.com/apikey")

    sys.exit(0 if all_ok else 1)


if __name__ == "__main__":
    main()
