#!/usr/bin/env python3
"""
Test script for multimodal AI backend service
Tests both text-only and image+text functionality
"""

import requests
import json
import time

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

def test_text_only():
    """Test text-only chat completion"""
    print("🧪 Testing text-only chat completion...")
    
    payload = {
        "model": "microsoft/DialoGPT-medium",
        "messages": [
            {"role": "user", "content": "Hello! How are you today?"}
        ],
        "max_tokens": 100,
        "temperature": 0.7
    }
    
    try:
        response = requests.post(f"{BASE_URL}/v1/chat/completions", json=payload, timeout=30)
        if response.status_code == 200:
            result = response.json()
            print(f"✅ Text-only response: {result['choices'][0]['message']['content']}")
            return True
        else:
            print(f"❌ Text-only failed: {response.status_code} - {response.text}")
            return False
    except Exception as e:
        print(f"❌ Text-only error: {e}")
        return False

def test_multimodal():
    """Test multimodal (image + text) chat completion"""
    print("🖼️ Testing multimodal chat completion...")
    
    payload = {
        "model": "unsloth/DeepSeek-R1-0528-Qwen3-8B-GGUF",
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/p-blog/candy.JPG"
                    },
                    {
                        "type": "text",
                        "text": "What animal is on the candy?"
                    }
                ]
            }
        ],
        "max_tokens": 150,
        "temperature": 0.7
    }
    
    try:
        response = requests.post(f"{BASE_URL}/v1/chat/completions", json=payload, timeout=60)
        if response.status_code == 200:
            result = response.json()
            print(f"✅ Multimodal response: {result['choices'][0]['message']['content']}")
            return True
        else:
            print(f"❌ Multimodal failed: {response.status_code} - {response.text}")
            return False
    except Exception as e:
        print(f"❌ Multimodal error: {e}")
        return False

def test_service_info():
    """Test service information endpoint"""
    print("ℹ️ Testing service information...")
    
    try:
        response = requests.get(f"{BASE_URL}/", timeout=10)
        if response.status_code == 200:
            result = response.json()
            print(f"✅ Service info: {result['message']}")
            return True
        else:
            print(f"❌ Service info failed: {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ Service info error: {e}")
        return False

def test_health():
    """Test health check endpoint"""
    print("🏥 Testing health check...")
    
    try:
        response = requests.get(f"{BASE_URL}/health", timeout=10)
        if response.status_code == 200:
            result = response.json()
            print(f"✅ Health: {result['status']} - Model: {result['model']}")
            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 main():
    """Run all tests"""
    print("🚀 Starting multimodal AI backend tests...\n")
    
    tests = [
        ("Service Info", test_service_info),
        ("Health Check", test_health),
        ("Text-only Chat", test_text_only),
        ("Multimodal Chat", test_multimodal),
    ]
    
    passed = 0
    total = len(tests)
    
    for test_name, test_func in tests:
        print(f"\n--- {test_name} ---")
        if test_func():
            passed += 1
        time.sleep(1)
    
    print(f"\n🎯 Test Results: {passed}/{total} tests passed")
    
    if passed == total:
        print("🎉 All tests passed! Multimodal AI backend is working correctly!")
    else:
        print("⚠️ Some tests failed. Check the output above for details.")

if __name__ == "__main__":
    main()