#!/usr/bin/env python3
"""
Test script to verify NVIDIA LLM Proxy connection for MathVista evaluation.
Uses Azure OpenAI format: /openai/deployments/{model}/chat/completions
"""

import os
import json
import requests

# NVIDIA LLM Proxy configuration
API_KEY = os.environ.get("NSTORAGE_API_KEY", "eyJhbGciOiJIUzI1NiJ9.eyJpZCI6IjdhNGZlOGYyLTQ1NjgtNGQzYS04OWYzLTE4ZDk3MmMwMmQ2NyIsInNlY3JldCI6Iit0aTViNVJ6S2F5dUNOQmxxdFNvdEdWUmJaSXJrRE03RDZLUjMrUjZZbDQ9In0.oyiXE5cg19mBweht3C-u5BAuRVJ5-QkcgqCt5W1sGGw")
API_VERSION = "2024-12-01-preview"

# Models to test
MODEL_NAMES = [
    "gpt-4o-20240806",
    "gpt-4o-mini-20240718",
    "gpt-4o-20241120",
    "gpt-4o-mini",
]

def get_api_url(model_name):
    """Azure format: /openai/deployments/{model}/chat/completions"""
    return f"https://llm-proxy.perflab.nvidia.com/openai/deployments/{model_name}/chat/completions?api-version={API_VERSION}"

def test_api(model_name, api_key):
    """Test API connection with given model."""
    url = get_api_url(model_name)
    
    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {api_key}',
    }
    
    payload = {
        "messages": [{"role": "user", "content": "Say hello in one word."}],
        "max_tokens": 10,
        "temperature": 0
    }
    
    print(f"\nTesting model: {model_name}")
    print(f"  URL: {url[:80]}...")
    
    try:
        response = requests.post(url, headers=headers, data=json.dumps(payload), timeout=30)
        
        if response.status_code == 200:
            resp_json = response.json()
            answer = resp_json['choices'][0]['message']['content']
            print(f"  ✅ SUCCESS! Answer: {answer}")
            return True, model_name
        elif response.status_code == 429:
            print(f"  ❌ 429 - No quota: {response.text[:100]}")
        elif response.status_code == 404:
            print(f"  ❌ 404 - Not found")
        else:
            print(f"  ❌ {response.status_code}: {response.text[:200]}")
        return False, None
            
    except Exception as e:
        print(f"  ❌ Error: {type(e).__name__}: {e}")
        return False, None

def main():
    print("="*60)
    print("Testing NVIDIA LLM Proxy (Azure OpenAI format)")
    print("="*60)
    print(f"API Version: {API_VERSION}")
    print(f"API Key: {API_KEY[:30]}...")
    
    working_models = []
    
    for model_name in MODEL_NAMES:
        success, model = test_api(model_name, API_KEY)
        if success:
            working_models.append(model)
    
    print("\n" + "="*60)
    if working_models:
        print("✅ WORKING MODELS FOUND:")
        for m in working_models:
            print(f"  - {m}")
        print(f"\nRecommended JUDGE_MODEL: {working_models[0]}")
    else:
        print("❌ No working models found.")
    print("="*60)

if __name__ == "__main__":
    main()
