#!/usr/bin/env python3
"""
Real API test for UltraGPT OpenRouter support
Tests OpenRouter with Claude models using minimal credit usage
"""

import os
import sys
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))

from dotenv import load_dotenv
from pydantic import BaseModel
from typing import List

# Load environment variables
load_dotenv()

from ultragpt import UltraGPT

class SimpleResponse(BaseModel):
    answer: str
    confidence: float
    reasoning: str

def test_claude_basic():
    """Test basic Claude functionality via OpenRouter"""
    print("� Testing Claude Basic Chat (via OpenRouter)...")
    
    try:
        ultragpt = UltraGPT(
            openrouter_api_key=os.getenv("OPENROUTER_API_KEY"),
            verbose=False
        )
        
        response, tokens, details = ultragpt.chat(
            model="claude:claude-3-haiku",
            messages=[{"role": "user", "content": "What is 2+2? Give a very brief answer."}],
            reasoning_pipeline=False,
            steps_pipeline=False,
            tools=[]
        )
        
        print(f"✓ Claude Response: {response}")
        print(f"✓ Tokens used: {tokens}")
        return True
        
    except Exception as e:
        print(f"✗ Claude basic test failed: {e}")
        import traceback
        traceback.print_exc()
        return False

def test_claude_sonnet_4():
    """Test Claude Sonnet 4 with 1M context"""
    print("\n🟠 Testing Claude Sonnet 4 (1M context via OpenRouter)...")
    
    try:
        ultragpt = UltraGPT(
            openrouter_api_key=os.getenv("OPENROUTER_API_KEY"),
            verbose=False
        )
        
        response, tokens, details = ultragpt.chat(
            model="claude:claude-sonnet-4",
            messages=[{"role": "user", "content": "What is 3+3? Give a very brief answer."}],
            reasoning_pipeline=False,
            steps_pipeline=False,
            tools=[]
        )
        
        print(f"✓ Claude Sonnet 4 Response: {response}")
        print(f"✓ Tokens used: {tokens}")
        return True
        
    except Exception as e:
        print(f"✗ Claude Sonnet 4 test failed: {e}")
        import traceback
        traceback.print_exc()
        return False

def test_claude_structured():
    """Test Claude structured output via OpenRouter"""
    print("\n� Testing Claude Structured Output (via OpenRouter)...")
    
    try:
        ultragpt = UltraGPT(
            openrouter_api_key=os.getenv("OPENROUTER_API_KEY"),
            verbose=False
        )
        
        response, tokens, details = ultragpt.chat(
            model="claude:claude-3-haiku",
            messages=[{"role": "user", "content": "Is the sky blue? Answer with confidence level."}],
            schema=SimpleResponse,
            reasoning_pipeline=False,
            steps_pipeline=False,
            tools=[]
        )
        
        print(f"✓ Claude Structured Response: {response}")
        print(f"✓ Tokens used: {tokens}")
        print(f"✓ Answer: {response.get('answer')}")
        print(f"✓ Confidence: {response.get('confidence')}")
        return True
        
    except Exception as e:
        print(f"✗ Claude structured test failed: {e}")
        import traceback
        traceback.print_exc()
        return False

def test_native_thinking():
    """Test native thinking support with Claude"""
    print("\n🧠 Testing Native Thinking Support (via OpenRouter)...")
    
    try:
        ultragpt = UltraGPT(
            openrouter_api_key=os.getenv("OPENROUTER_API_KEY"),
            verbose=True  # Set verbose to see native thinking detection message
        )
        
        # Test with reasoning_pipeline=True (should detect native thinking and use it)
        response, tokens, details = ultragpt.chat(
            model="claude:claude-sonnet-4",
            messages=[{"role": "user", "content": "What is the square root of 144?"}],
            reasoning_pipeline=True,  # Should trigger native thinking detection
            steps_pipeline=False,
            tools=[]
        )
        
        print(f"✓ Native Thinking Response: {response}")
        print(f"✓ Tokens used: {tokens}")
        print(f"✓ Details keys: {details.keys()}")
        return True
        
    except Exception as e:
        print(f"✗ Native thinking test failed: {e}")
        import traceback
        traceback.print_exc()
        return False

def test_message_conversion():
    """Test Claude message conversion with real API via OpenRouter"""
    print("\n🟠 Testing Claude Message Conversion (via OpenRouter)...")
    
    try:
        ultragpt = UltraGPT(
            openrouter_api_key=os.getenv("OPENROUTER_API_KEY"),
            verbose=False
        )
        
        # Test with system message
        messages = [
            {"role": "system", "content": "You are a helpful assistant. Be very brief."},
            {"role": "user", "content": "Hello, how are you?"}
        ]
        
        response, tokens, details = ultragpt.chat(
            model="claude:claude-3-haiku",
            messages=messages,
            reasoning_pipeline=False,
            steps_pipeline=False,
            tools=[]
        )
        
        print(f"✓ Claude System Message Response: {response}")
        print(f"✓ Tokens used: {tokens}")
        return True
        
    except Exception as e:
        print(f"✗ Claude message conversion test failed: {e}")
        import traceback
        traceback.print_exc()
        return False

def main():
    """Run minimal real API tests with OpenRouter"""
    print("🚀 UltraGPT Real API Tests (OpenRouter)")
    print("=" * 50)
    print("Running minimal tests to verify OpenRouter functionality...")
    print("(Designed to use minimal API credits)")
    print("=" * 50)
    
    results = []
    total_tests = 0
    
    # Check API key
    if not os.getenv("OPENROUTER_API_KEY"):
        print("❌ OPENROUTER_API_KEY not set!")
        print("Please set OPENROUTER_API_KEY in your .env file")
        return
    
    print("✓ OPENROUTER_API_KEY found")
    
    # Run tests
    total_tests = 5
    results.append(("Claude Basic (Haiku)", test_claude_basic()))
    results.append(("Claude Sonnet 4 (1M)", test_claude_sonnet_4()))
    results.append(("Claude Structured", test_claude_structured()))
    results.append(("Native Thinking", test_native_thinking()))
    results.append(("Claude Messages", test_message_conversion()))
    
    # Summary
    print("\n" + "=" * 50)
    print("📊 TEST RESULTS SUMMARY")
    print("=" * 50)
    
    passed = 0
    for test_name, result in results:
        status = "✅ PASS" if result else "❌ FAIL"
        print(f"{test_name:25} {status}")
        if result:
            passed += 1
    
    print(f"\n🎯 Overall: {passed}/{total_tests} tests passed")
    
    if passed == total_tests:
        print("🎉 All tests passed! OpenRouter is working correctly.")
    else:
        print("⚠️  Some tests failed. Check the output above for details.")
    
    print("\n💡 Features tested:")
    print("- ✓ Basic chat with Claude models")
    print("- ✓ Claude Sonnet 4 with 1M extended context")
    print("- ✓ Structured output (Pydantic schemas)")
    print("- ✓ Native thinking support (deepthink=True)")
    print("- ✓ System message handling")
    print("\n🔥 OpenRouter is now ready for production use!")

if __name__ == "__main__":
    main()
