import os
import requests
import json
from dotenv import load_dotenv

# Load environment variables from the project root
load_dotenv(override=True)

HF_TOKEN = os.getenv("HF_TOKEN")
HF_ROUTER_URL = "https://router.huggingface.co/v1"
HF_INFERENCE_URL = "https://api-inference.huggingface.co/models"

def test_model_router(model_name):
    print(f"\n--- Testing via ROUTER: {model_name} ---")
    headers = {"Authorization": f"Bearer {HF_TOKEN}", "Content-Type": "application/json"}
    payload = {
        "model": model_name,
        "messages": [{"role": "user", "content": "RPG character prompt: Wizard"}],
        "max_tokens": 50
    }
    try:
        response = requests.post(f"{HF_ROUTER_URL}/chat/completions", headers=headers, json=payload, timeout=10)
        if response.status_code == 200:
            data = response.json()
            print(f"✅ SUCCESS Raw Result: {json.dumps(data, indent=2)}")
        else:
            print(f"⌠FAILED ({response.status_code}): {response.text}")
    except Exception as e:
        print(f"âš ï¸ ERROR: {e}")

def test_model_inference_api(model_name):
    print(f"\n--- Testing via DIRECT INFERENCE API: {model_name} ---")
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    payload = {
        "inputs": "RPG character prompt: Wizard",
        "parameters": {"max_new_tokens": 50}
    }
    try:
        response = requests.post(f"{HF_INFERENCE_URL}/{model_name}", headers=headers, json=payload, timeout=10)
        if response.status_code == 200:
            print(f"✅ SUCCESS: {response.json()}")
        else:
            print(f"⌠FAILED ({response.status_code}): {response.text}")
    except Exception as e:
        print(f"âš ï¸ ERROR: {e}")

if __name__ == "__main__":
    if not HF_TOKEN:
        print("Error: HF_TOKEN not found in .env")
        exit(1)
        
    # Candidates for Gemma 4 (April 2026)
    candidates = [
        "google/gemma-4-31b-it",
        "google/gemma-4-31B-it",
        "google/gemma-4-26b-a4b-it",
        "google/gemma-g4-31b-it" # Alternative naming
    ]
    
    for model in candidates:
        test_model_router(model)
        test_model_inference_api(model)