#!/usr/bin/env python3
"""
🧪 Script de Prueba para Gráficos Vectoriales SVG
Prueba la funcionalidad de generación de gráficos vectoriales con SVGDreamer
"""

import os
import sys
import time
import tempfile
import requests

def test_svgdreamer_basic():
    """Prueba básica de SVGDreamer"""
    print("🎨 Probando SVGDreamer - Generación básica...")
    
    try:
        # Configurar API
        API_URL = "https://api-inference.huggingface.co/models/jree423/svgdreamer"
        hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
        headers = {"Authorization": f"Bearer {hf_token}"} if hf_token else {}
        
        # Payload básico
        payload = {
            "inputs": "a simple house icon",
            "parameters": {
                "n_particle": 1,
                "num_iter": 100,  # Bajo para prueba rápida
                "guidance_scale": 7.5,
                "style": "iconography",
                "width": 224,
                "height": 224,
                "seed": 42
            }
        }
        
        print(f"📦 Enviando payload: {payload}")
        
        # Realizar request
        start_time = time.time()
        response = requests.post(API_URL, headers=headers, json=payload)
        generation_time = time.time() - start_time
        
        if response.status_code != 200:
            raise Exception(f"Error en API: {response.status_code} - {response.text}")
        
        result = response.json()
        print(f"✅ Respuesta recibida en {generation_time:.2f}s")
        print(f"📊 Tipo de respuesta: {type(result)}")
        
        # Procesar respuesta
        if isinstance(result, dict) and 'generated_text' in result:
            svg_content = result['generated_text']
        elif isinstance(result, list):
            svg_content = result
        else:
            svg_content = result
        
        print(f"📄 Contenido SVG recibido: {len(str(svg_content))} caracteres")
        
        # Guardar archivo de prueba
        with tempfile.NamedTemporaryFile(suffix='.svg', delete=False, mode='w', encoding='utf-8') as tmp_file:
            tmp_file.write(str(svg_content))
            test_file = tmp_file.name
        
        print(f"💾 Archivo de prueba guardado: {test_file}")
        
        # Verificar que el archivo contiene SVG válido
        with open(test_file, 'r', encoding='utf-8') as f:
            content = f.read()
            if '<svg' in content.lower():
                print("✅ Archivo SVG válido generado")
            else:
                print("⚠️ El archivo no parece contener SVG válido")
        
        # Limpiar archivo temporal
        os.unlink(test_file)
        
        return True
        
    except Exception as e:
        print(f"❌ Error en prueba básica: {e}")
        return False

def test_svgdreamer_multiple_styles():
    """Prueba SVGDreamer con diferentes estilos"""
    print("\n🎨 Probando SVGDreamer - Múltiples estilos...")
    
    styles = ["iconography", "pixel_art", "sketch", "painting"]
    prompt = "a friendly robot character"
    
    try:
        API_URL = "https://api-inference.huggingface.co/models/jree423/svgdreamer"
        hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
        headers = {"Authorization": f"Bearer {hf_token}"} if hf_token else {}
        
        for style in styles:
            print(f"🎯 Probando estilo: {style}")
            
            payload = {
                "inputs": prompt,
                "parameters": {
                    "n_particle": 1,
                    "num_iter": 50,  # Muy bajo para pruebas rápidas
                    "guidance_scale": 7.5,
                    "style": style,
                    "width": 128,  # Más pequeño para velocidad
                    "height": 128,
                    "seed": 42
                }
            }
            
            start_time = time.time()
            response = requests.post(API_URL, headers=headers, json=payload)
            generation_time = time.time() - start_time
            
            if response.status_code == 200:
                print(f"✅ Estilo {style} completado en {generation_time:.2f}s")
            else:
                print(f"❌ Error con estilo {style}: {response.status_code}")
            
            # Pausa entre requests para no sobrecargar
            time.sleep(1)
        
        print("✅ Todos los estilos probados exitosamente")
        return True
        
    except Exception as e:
        print(f"❌ Error en prueba de estilos: {e}")
        return False

def test_svgdreamer_multiple_particles():
    """Prueba SVGDreamer con múltiples partículas"""
    print("\n🎨 Probando SVGDreamer - Múltiples partículas...")
    
    try:
        API_URL = "https://api-inference.huggingface.co/models/jree423/svgdreamer"
        hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
        headers = {"Authorization": f"Bearer {hf_token}"} if hf_token else {}
        
        payload = {
            "inputs": "geometric patterns in bright colors",
            "parameters": {
                "n_particle": 3,  # Múltiples partículas
                "num_iter": 50,   # Bajo para velocidad
                "guidance_scale": 7.5,
                "style": "iconography",
                "width": 128,
                "height": 128,
                "seed": 42
            }
        }
        
        print(f"📦 Enviando payload con 3 partículas...")
        
        start_time = time.time()
        response = requests.post(API_URL, headers=headers, json=payload)
        generation_time = time.time() - start_time
        
        if response.status_code != 200:
            raise Exception(f"Error en API: {response.status_code} - {response.text}")
        
        result = response.json()
        print(f"✅ Respuesta recibida en {generation_time:.2f}s")
        
        # Procesar respuesta
        if isinstance(result, dict) and 'generated_text' in result:
            svg_content = result['generated_text']
        elif isinstance(result, list):
            svg_content = result
        else:
            svg_content = result
        
        # Verificar si es una lista de partículas
        if isinstance(svg_content, list):
            print(f"✅ Generadas {len(svg_content)} partículas")
            for i, particle in enumerate(svg_content):
                if isinstance(particle, dict) and 'svg' in particle:
                    print(f"  📄 Partícula {i+1}: {len(particle['svg'])} caracteres")
                else:
                    print(f"  📄 Partícula {i+1}: {len(str(particle))} caracteres")
        else:
            print(f"📄 Respuesta única: {len(str(svg_content))} caracteres")
        
        return True
        
    except Exception as e:
        print(f"❌ Error en prueba de partículas: {e}")
        return False

def test_error_handling():
    """Prueba el manejo de errores"""
    print("\n🎨 Probando manejo de errores...")
    
    try:
        API_URL = "https://api-inference.huggingface.co/models/jree423/svgdreamer"
        hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
        headers = {"Authorization": f"Bearer {hf_token}"} if hf_token else {}
        
        # Payload con parámetros inválidos
        payload = {
            "inputs": "test",
            "parameters": {
                "n_particle": 20,  # Demasiado alto
                "num_iter": 2000,  # Demasiado alto
                "guidance_scale": 50.0,  # Demasiado alto
                "style": "invalid_style",  # Estilo inválido
                "width": 1000,  # Demasiado grande
                "height": 1000,
                "seed": 42
            }
        }
        
        print("🧪 Probando parámetros extremos...")
        
        try:
            response = requests.post(API_URL, headers=headers, json=payload)
            if response.status_code == 200:
                print("⚠️ Request completado (esperaba error)")
            else:
                print(f"✅ Error capturado correctamente: {response.status_code}")
        except Exception as e:
            print(f"✅ Error capturado correctamente: {type(e).__name__}")
        
        return True
        
    except Exception as e:
        print(f"❌ Error en prueba de manejo de errores: {e}")
        return False

def main():
    """Función principal de pruebas"""
    print("🧪 Iniciando pruebas de gráficos vectoriales SVG...")
    print("=" * 60)
    
    # Configurar token si está disponible
    hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
    if hf_token:
        print(f"🔑 Token detectado: {hf_token[:10]}...")
    else:
        print("⚠️ No se detectó HF_TOKEN - algunas funcionalidades pueden estar limitadas")
    
    # Ejecutar pruebas
    tests = [
        ("Prueba básica", test_svgdreamer_basic),
        ("Múltiples estilos", test_svgdreamer_multiple_styles),
        ("Múltiples partículas", test_svgdreamer_multiple_particles),
        ("Manejo de errores", test_error_handling)
    ]
    
    results = []
    
    for test_name, test_func in tests:
        print(f"\n{'='*20} {test_name} {'='*20}")
        try:
            result = test_func()
            results.append((test_name, result))
        except Exception as e:
            print(f"❌ Error inesperado en {test_name}: {e}")
            results.append((test_name, False))
    
    # Resumen de resultados
    print(f"\n{'='*60}")
    print("📊 RESUMEN DE PRUEBAS")
    print("=" * 60)
    
    passed = 0
    total = len(results)
    
    for test_name, result in results:
        status = "✅ PASÓ" if result else "❌ FALLÓ"
        print(f"{test_name}: {status}")
        if result:
            passed += 1
    
    print(f"\n🎯 Resultado: {passed}/{total} pruebas pasaron")
    
    if passed == total:
        print("🎉 ¡Todas las pruebas pasaron! La funcionalidad está lista.")
    else:
        print("⚠️ Algunas pruebas fallaron. Revisar configuración.")
    
    return passed == total

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)