#!/usr/bin/env python3
"""
Test script for the Mistral Nemo Backend Service
"""

import requests
import json
import time

# Service configuration
BASE_URL = "http://localhost:8001"

def test_health():
    """Test the health endpoint"""
    print("🏥 Testing health endpoint...")
    try:
        response = requests.get(f"{BASE_URL}/health", timeout=5)
        if response.status_code == 200:
            print(f"✅ Health check passed: {response.json()}")
            return True
        else:
            print(f"❌ Health check failed: {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ Health check error: {e}")
        return False

def test_root():
    """Test the root endpoint"""
    print("🏠 Testing root endpoint...")
    try:
        response = requests.get(f"{BASE_URL}/", timeout=5)
        if response.status_code == 200:
            data = response.json()
            print(f"✅ Root endpoint: {data}")
            return True
        else:
            print(f"❌ Root endpoint failed: {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ Root endpoint error: {e}")
        return False

def test_models():
    """Test the models endpoint"""
    print("📋 Testing models endpoint...")
    try:
        response = requests.get(f"{BASE_URL}/v1/models", timeout=5)
        if response.status_code == 200:
            data = response.json()
            print(f"✅ Available models: {[model['id'] for model in data['data']]}")
            return True
        else:
            print(f"❌ Models endpoint failed: {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ Models endpoint error: {e}")
        return False

def test_chat_completion():
    """Test a simple chat completion"""
    print("💬 Testing chat completion...")
    try:
        payload = {
            "model": "unsloth/Mistral-Nemo-Instruct-2407",
            "messages": [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": "Hello! Tell me a fun fact about AI."}
            ],
            "max_tokens": 100,
            "temperature": 0.7
        }
        
        response = requests.post(f"{BASE_URL}/v1/chat/completions", 
                               json=payload, timeout=30)
        
        if response.status_code == 200:
            data = response.json()
            message = data["choices"][0]["message"]["content"]
            print(f"✅ Chat completion successful!")
            print(f"🤖 Assistant: {message}")
            return True
        else:
            print(f"❌ Chat completion failed: {response.status_code}")
            print(f"Response: {response.text}")
            return False
    except Exception as e:
        print(f"❌ Chat completion error: {e}")
        return False

def wait_for_service():
    """Wait for the service to be ready"""
    print("⏳ Waiting for service to be ready...")
    max_attempts = 60  # Wait up to 5 minutes
    
    for attempt in range(max_attempts):
        try:
            response = requests.get(f"{BASE_URL}/health", timeout=5)
            if response.status_code == 200:
                print(f"✅ Service is ready after {attempt * 5} seconds!")
                return True
        except:
            pass
        
        if attempt < max_attempts - 1:
            print(f"⏳ Attempt {attempt + 1}/{max_attempts} - waiting 5 seconds...")
            time.sleep(5)
    
    print("❌ Service did not become ready within the timeout period")
    return False

def main():
    """Run all tests"""
    print("🚀 Testing Mistral Nemo Backend Service")
    print("=" * 50)
    
    # Wait for service to be ready
    if not wait_for_service():
        print("❌ Service is not ready. Exiting.")
        return
    
    # Run tests
    tests = [test_root, test_health, test_models, test_chat_completion]
    passed = 0
    
    for test in tests:
        try:
            if test():
                passed += 1
            print()
        except Exception as e:
            print(f"❌ Test failed with exception: {e}")
            print()
    
    print("=" * 50)
    print(f"📊 Test Results: {passed}/{len(tests)} tests passed")
    
    if passed == len(tests):
        print("🎉 All tests passed! Your Mistral Nemo service is working perfectly!")
    else:
        print("⚠️ Some tests failed. Check the logs above for details.")

if __name__ == "__main__":
    main()