import os
import sys

# Color definitions
GREEN = "\033[0;32m"
RED = "\033[0;31m"
YELLOW = "\033[0;33m"
BLUE = "\033[0;34m"
NC = "\033[0m"

print(f"{BLUE}==================================================={NC}")
print(f"{BLUE}   Gemini API Key Connection Diagnostics           {NC}")
print(f"{BLUE}==================================================={NC}\n")

# Load key
api_key = os.getenv("GEMINI_API_KEY")
if not api_key:
    if os.path.exists(".env"):
        try:
            with open(".env", "r", encoding="utf-8") as f:
                for line in f:
                    if line.strip().startswith("GEMINI_API_KEY="):
                        api_key = line.strip().split("=", 1)[1].strip("'\" ")
                        break
        except Exception as e:
            print(f"{RED}[!] Error reading .env file: {e}{NC}")
    else:
        print(f"{YELLOW}[!] Warning: .env file not found.{NC}")

if not api_key:
    print(f"{RED}[❌] GEMINI_API_KEY is not configured!{NC}")
    print("Please set it in your .env file or export it in your environment.")
    sys.exit(1)

print(f"[*] Found API Key: {api_key[:12]}...{api_key[-5:] if len(api_key) > 15 else ''}")

try:
    from google import genai
except ImportError:
    print(f"{YELLOW}[!] 'google-genai' library not found in global python. Trying to import from virtual env...{NC}")
    venv_path = os.path.join(os.getcwd(), ".venv", "lib", "python3.12", "site-packages")
    if os.path.exists(venv_path):
        sys.path.insert(0, venv_path)
    try:
        from google import genai
    except ImportError:
        print(f"{RED}[❌] Failed to import google-genai. Please run using: .venv/bin/python test_gemini.py{NC}")
        sys.exit(1)

client = genai.Client(api_key=api_key)

test_models = [
    'gemini-2.5-flash-lite',
    'gemini-3.1-flash-lite',
    'gemini-2.5-flash',
    'gemini-3.5-flash',
    'gemini-flash-lite-latest',
]

working_count = 0

for model in test_models:
    print(f"[*] Testing model: {model} ... ", end="")
    sys.stdout.flush()
    try:
        response = client.models.generate_content(
            model=model,
            contents="Say 'Hello' in Vietnamese in 3 words"
        )
        print(f"{GREEN}[✅ WORKING]{NC} -> Response: \"{response.text.strip()}\"")
        working_count += 1
    except Exception as e:
        err_msg = str(e).strip()
        first_line = err_msg.split('\n')[0]
        # Clean error message
        if "429" in first_line or "RESOURCE_EXHAUSTED" in first_line:
            print(f"{RED}[❌ LIMIT/QUOTA EXHAUSTED]{NC} (Rate limit exceeded for this model)")
        elif "503" in first_line or "UNAVAILABLE" in first_line:
            print(f"{RED}[❌ TEMPORARILY UNAVAILABLE]{NC} (Model under heavy load)")
        else:
            print(f"{RED}[❌ ERROR]{NC} -> {first_line[:120]}")

print(f"\n{BLUE}==================================================={NC}")
if working_count > 0:
    print(f"{GREEN}[✓] Diagnostics Completed. {working_count}/{len(test_models)} models are working successfully!{NC}")
    print(f"The app fallback mechanism will automatically use the working models.")
else:
    print(f"{RED}[❌] All models failed. Please verify if your API Key is correct or has active credit/billing.{NC}")
print(f"{BLUE}==================================================={NC}")