"""
UI-TARS API Test Client (Optimized Version)
===========================================
اختبار سريع للـ API المحسّن
"""

import requests
import base64
import time
from io import BytesIO
from PIL import Image

# Configuration
API_URL = "http://localhost:7860"  # غيّره لـ Space URL الخاص بك

def create_test_image():
    """إنشاء صورة اختبار"""
    img = Image.new('RGB', (1920, 1080), color='white')
    # رسم مربع في المنتصف
    from PIL import ImageDraw
    draw = ImageDraw.Draw(img)
    draw.rectangle([900, 500, 1020, 580], outline='red', width=3)
    draw.text((910, 530), "Button", fill='red')
    
    buffer = BytesIO()
    img.save(buffer, format='PNG')
    return base64.b64encode(buffer.getvalue()).decode()

def test_health():
    """اختبار endpoint الصحة"""
    print("\n" + "="*60)
    print("🔍 Testing Health Endpoint")
    print("="*60)
    
    try:
        response = requests.get(f"{API_URL}/health", timeout=10)
        print(f"✅ Status Code: {response.status_code}")
        data = response.json()
        print(f"📊 Response:")
        for key, value in data.items():
            print(f"   {key}: {value}")
        return True
    except Exception as e:
        print(f"❌ Error: {e}")
        return False

def test_model_info():
    """اختبار معلومات النموذج"""
    print("\n" + "="*60)
    print("📋 Testing Model Info Endpoint")
    print("="*60)
    
    try:
        response = requests.get(f"{API_URL}/model/info", timeout=10)
        print(f"✅ Status Code: {response.status_code}")
        data = response.json()
        print(f"📊 Model Info:")
        print(f"   Name: {data.get('model_name')}")
        print(f"   API Type: {data.get('api_type')}")
        print(f"   Capabilities: {', '.join(data.get('capabilities', []))}")
        return True
    except Exception as e:
        print(f"❌ Error: {e}")
        return False

def test_inference_simple():
    """اختبار استدلال بسيط (بدون صورة)"""
    print("\n" + "="*60)
    print("🤖 Testing Simple Inference (No Image)")
    print("="*60)
    
    try:
        payload = {
            "instruction": "Click on the start button",
            "system_prompt_type": "computer"
        }
        
        print("⏳ Sending request...")
        response = requests.post(
            f"{API_URL}/v1/inference",
            json=payload,
            timeout=60
        )
        
        print(f"✅ Status Code: {response.status_code}")
        
        if response.status_code == 200:
            data = response.json()
            print(f"💭 Thought: {data.get('thought', 'N/A')[:100]}...")
            print(f"⚡ Action: {data.get('action', 'N/A')}")
            if data.get('coordinates'):
                print(f"📍 Coordinates: {data['coordinates']}")
            return True
        else:
            print(f"❌ Error Response: {response.text[:200]}")
            return False
            
    except Exception as e:
        print(f"❌ Error: {e}")
        return False

def test_inference_with_image():
    """اختبار استدلال مع صورة"""
    print("\n" + "="*60)
    print("🖼️  Testing Inference With Image")
    print("="*60)
    
    try:
        image_b64 = create_test_image()
        print(f"✅ Test image created (size: {len(image_b64)} chars)")
        
        payload = {
            "instruction": "Click on the red button in the center",
            "image": image_b64,
            "system_prompt_type": "computer",
            "max_tokens": 512
        }
        
        print("⏳ Sending request...")
        response = requests.post(
            f"{API_URL}/v1/inference",
            json=payload,
            timeout=60
        )
        
        print(f"✅ Status Code: {response.status_code}")
        
        if response.status_code == 200:
            data = response.json()
            print(f"💭 Thought: {data.get('thought', 'N/A')[:100]}...")
            print(f"⚡ Action: {data.get('action', 'N/A')}")
            if data.get('coordinates'):
                coords = data['coordinates']
                print(f"📍 Coordinates: x={coords['x']}, y={coords['y']}")
            return True
        else:
            print(f"❌ Error Response: {response.text[:200]}")
            # إذا كان النموذج يحمّل، انتظر وحاول مرة أخرى
            if "loading" in response.text.lower():
                print("⏳ Model is loading... waiting 15 seconds...")
                time.sleep(15)
                return test_inference_with_image()  # إعادة المحاولة
            return False
            
    except Exception as e:
        print(f"❌ Error: {e}")
        return False

def test_chat_completion():
    """اختبار OpenAI-compatible endpoint"""
    print("\n" + "="*60)
    print("💬 Testing Chat Completion Endpoint")
    print("="*60)
    
    try:
        image_b64 = create_test_image()
        
        payload = {
            "model": "ui-tars-1.5-7b",
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": "Click on the button"
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{image_b64}"
                            }
                        }
                    ]
                }
            ],
            "max_tokens": 512
        }
        
        print("⏳ Sending request...")
        response = requests.post(
            f"{API_URL}/v1/chat/completions",
            json=payload,
            timeout=60
        )
        
        print(f"✅ Status Code: {response.status_code}")
        
        if response.status_code == 200:
            data = response.json()
            content = data["choices"][0]["message"]["content"]
            print(f"💬 Response: {content[:150]}...")
            return True
        else:
            print(f"❌ Error Response: {response.text[:200]}")
            return False
            
    except Exception as e:
        print(f"❌ Error: {e}")
        return False

def run_all_tests():
    """تشغيل جميع الاختبارات"""
    print("\n" + "="*60)
    print("🚀 UI-TARS API Test Suite (Optimized)")
    print("="*60)
    print(f"🔗 Testing API: {API_URL}")
    
    results = {
        "Health Check": test_health(),
        "Model Info": test_model_info(),
        "Simple Inference": test_inference_simple(),
        "Inference with Image": test_inference_with_image(),
        "Chat Completion": test_chat_completion()
    }
    
    # النتائج النهائية
    print("\n" + "="*60)
    print("📊 Test Results Summary")
    print("="*60)
    
    for test_name, passed in results.items():
        status = "✅ PASSED" if passed else "❌ FAILED"
        print(f"{test_name:.<40} {status}")
    
    total = len(results)
    passed = sum(results.values())
    
    print("="*60)
    print(f"Total: {passed}/{total} tests passed ({passed/total*100:.1f}%)")
    print("="*60)
    
    return passed == total

if __name__ == "__main__":
    # يمكنك تغيير API_URL هنا
    # API_URL = "https://your-space.hf.space"
    
    success = run_all_tests()
    
    if success:
        print("\n🎉 All tests passed! API is working perfectly.")
    else:
        print("\n⚠️  Some tests failed. Check the errors above.")
    
    exit(0 if success else 1)