#!/usr/bin/env python3
"""
Test script to verify OpenRouter API key and translation functionality
"""

import os
import asyncio
import aiohttp
from translator import DocumentTranslator

async def test_api_key():
    """Test if the API key is working"""
    print("🔑 Testing OpenRouter API key...")
    
    api_key = os.getenv("OPENROUTER_API_KEY")
    if not api_key:
        print("❌ OPENROUTER_API_KEY environment variable not set!")
        print("Please set it with: set OPENROUTER_API_KEY=your_key_here")
        return False
    
    print(f"✅ API key found: {api_key[:10]}...")
    
    # Test API connection
    try:
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://huggingface.co",
            "X-Title": "Document Translator"
        }
        
        async with aiohttp.ClientSession() as session:
            async with session.get(
                "https://openrouter.ai/api/v1/models",
                headers=headers
            ) as response:
                if response.status == 200:
                    print("✅ API connection successful!")
                    return True
                elif response.status == 429:
                    error_text = await response.text()
                    print(f"⚠️ Rate limit exceeded: {error_text}")
                    print("This is normal for free models. Try again later or use a different model.")
                    return True  # API key is valid, just rate limited
                else:
                    print(f"❌ API connection failed: {response.status}")
                    error_text = await response.text()
                    print(f"Error: {error_text}")
                    return False
    except Exception as e:
        print(f"❌ API test failed: {e}")
        return False

async def test_translation():
    """Test basic translation functionality"""
    print("\n📝 Testing translation functionality...")
    
    translator = DocumentTranslator()
    
    if not translator.is_ready():
        print("❌ Translator not ready - API key issue")
        return False
    
    try:
        # Get available models first
        models = await translator.get_available_models()
        if not models:
            print("❌ No models available")
            return False
        
        selected_model = models[0]["id"]
        print(f"Using model: {selected_model}")
        
        # Test simple translation
        test_text = "Hello, this is a test document."
        print(f"Original text: {test_text}")
        
        translated = await translator.translate_text(
            text=test_text,
            model=selected_model,
            source_lang="en",
            target_lang="ar"
        )
        
        print(f"Translated text: {translated}")
        
        if translated != test_text:
            print("✅ Translation working correctly!")
            return True
        else:
            print("❌ Translation returned original text - may indicate an issue")
            return False
            
    except Exception as e:
        print(f"❌ Translation test failed: {e}")
        return False

async def test_specific_model(model_id: str):
    """Test a specific model for translation"""
    print(f"\n🧪 Testing model: {model_id}")
    
    api_key = os.getenv("OPENROUTER_API_KEY")
    if not api_key:
        print("❌ OPENROUTER_API_KEY not set")
        return False
        
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://huggingface.co",
        "X-Title": "Document Translator"
    }
    
    test_payload = {
        "model": model_id,
        "messages": [
            {"role": "system", "content": "You are a professional translator."},
            {"role": "user", "content": "Translate 'Hello world' to Arabic"}
        ],
        "max_tokens": 50,
        "temperature": 0.1
    }
    
    try:
        async with aiohttp.ClientSession() as session:
            async with session.post(
                "https://openrouter.ai/api/v1/chat/completions",
                headers=headers,
                json=test_payload
            ) as response:
                if response.status == 200:
                    data = await response.json()
                    result = data["choices"][0]["message"]["content"]
                    print(f"✅ Model works! Translation: {result}")
                    return True
                elif response.status == 429:
                    error_text = await response.text()
                    print(f"⚠️ Model rate limited: {error_text}")
                    print("Try again later or use a different model.")
                    return True  # Model exists, just rate limited
                else:
                    error_text = await response.text()
                    print(f"❌ Model test failed: {response.status} - {error_text}")
                    return False
    except Exception as e:
        print(f"❌ Test error: {e}")
        return False

async def main():
    """Run all tests"""
    print("🧪 Testing Document Translator Setup\n")
    
    # Test API key
    api_ok = await test_api_key()
    
    if api_ok:
        # Test translation
        translation_ok = await test_translation()
        
        if translation_ok:
            print("\n🎉 All tests passed! The translator should work correctly.")
        else:
            print("\n⚠️ Translation test failed. Check the logs for details.")
    else:
        print("\n❌ API key test failed. Please check your OPENROUTER_API_KEY.")
    
    print("\n📋 Next steps:")
    print("1. Make sure OPENROUTER_API_KEY is set correctly")
    print("2. Upload a PDF or DOCX file to test the full workflow")
    print("3. Check the translation.log file for detailed logs")

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