#!/usr/bin/env python3
"""Test LiteLLM integration with OpenRouter."""

import asyncio
import os
import sys

from dotenv import load_dotenv
load_dotenv()

from litellm import acompletion


async def test_litellm_openrouter():
    """Test OpenRouter through LiteLLM."""
    api_key = os.environ.get("OPENROUTER_API_KEY")
    
    if not api_key:
        print("ERROR: OPENROUTER_API_KEY not set!")
        return False
    
    print(f"Testing LiteLLM with OpenRouter...")
    print(f"API Key: {api_key[:25]}...")
    
    # Test models
    models = [
        "openrouter/meta-llama/llama-3.3-70b-instruct",
        "openrouter/meta-llama/llama-3.1-8b-instruct",
    ]
    
    api_base = "https://openrouter.ai/api/v1"
    extra_headers = {
        "HTTP-Referer": os.environ.get("OPENROUTER_REFERER", "https://localhost"),
        "X-Title": os.environ.get("OPENROUTER_APP_TITLE", "HF Agent Test"),
    }
    
    for model in models:
        print(f"\n{'='*60}")
        print(f"Testing model: {model}")
        print(f"{'='*60}")
        
        try:
            # Test non-streaming
            print("Testing non-streaming...")
            response = await acompletion(
                model=model,
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": "Say 'LiteLLM works!' and nothing else."}
                ],
                max_tokens=50,
                temperature=0.1,
                api_key=api_key,
                api_base=api_base,
                extra_headers=extra_headers,
            )
            
            content = response.choices[0].message.content
            print(f"✅ Non-streaming SUCCESS! Response: {content}")
            
            # Test streaming
            print("\nTesting streaming...")
            stream_response = await acompletion(
                model=model,
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": "Say 'Streaming works!' and nothing else."}
                ],
                max_tokens=50,
                temperature=0.1,
                api_key=api_key,
                api_base=api_base,
                extra_headers=extra_headers,
                stream=True,
                stream_options={"include_usage": True},
            )
            
            streamed_content = ""
            async for chunk in stream_response:
                if chunk.choices and chunk.choices[0].delta.content:
                    streamed_content += chunk.choices[0].delta.content
            
            print(f"✅ Streaming SUCCESS! Response: {streamed_content}")
            
            return True
            
        except Exception as e:
            print(f"❌ FAILED: {type(e).__name__}: {e}")
            import traceback
            traceback.print_exc()
            continue
    
    return False


async def main():
    print("="*60)
    print("LiteLLM + OpenRouter Integration Test")
    print("="*60)
    
    success = await test_litellm_openrouter()
    
    if success:
        print("\n🎉 LiteLLM + OpenRouter integration is working!")
        return 0
    else:
        print("\n⚠️ LiteLLM + OpenRouter test failed!")
        return 1


if __name__ == "__main__":
    exit_code = asyncio.run(main())
    sys.exit(exit_code)