import gradio as gr
import torch
from PIL import Image
import numpy as np
from clip_interrogator import Config, Interrogator, LabelTable, load_list
import logging
import os
import warnings
from datetime import datetime
import json
import gc

# Suprimir warnings específicos
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning)
os.environ["TOKENIZERS_PARALLELISM"] = "false"

# Configurar logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Detectar dispositivo disponible
def get_device():
    if torch.cuda.is_available():
        return "cuda"
    elif torch.backends.mps.is_available():
        return "mps"
    else:
        return "cpu"

DEVICE = get_device()
logger.info(f"🖥️ Usando dispositivo: {DEVICE}")

# Configuración optimizada
CLIP_MODELS = {
    "general": "ViT-L-14/openai",
    "stable_diffusion": "ViT-L-14/openai",
    "midjourney": "ViT-H-14/laion2b_s32b_b79k",
    "flux": "ViT-L-14/openai"
}

INTERROGATION_MODES = {
    "fast": "⚡ Rápido (30 seg)",
    "classic": "⚖️ Clásico (1 min)", 
    "best": "⭐ Mejor (2 min)"
}

class OptimizedImagePromptGenerator:
    def __init__(self):
        self.interrogator = None
        self.usage_count = 0
        self.device = DEVICE
        self.is_initialized = False
        logger.info("🚀 Inicializando generador optimizado...")
    
    def initialize_model(self, progress_callback=None):
        """Inicialización lazy del modelo"""
        if self.is_initialized:
            return True
            
        try:
            if progress_callback:
                progress_callback("🔄 Configurando modelo CLIP...")
            
            # Configuración optimizada según dispositivo
            config = Config(
                clip_model_name="ViT-L-14/openai",
                download_cache=True,
                chunk_size=1024 if self.device == "cpu" else 2048,
                quiet=True,
                device=self.device
            )
            
            if progress_callback:
                progress_callback("📥 Descargando modelos (primera vez)...")
            
            self.interrogator = Interrogator(config)
            
            if progress_callback:
                progress_callback("✅ Modelo inicializado correctamente")
            
            self.is_initialized = True
            logger.info("✅ Modelo CLIP inicializado correctamente")
            
            # Limpiar memoria
            if self.device == "cpu":
                gc.collect()
            else:
                torch.cuda.empty_cache()
                
            return True
            
        except Exception as e:
            logger.error(f"❌ Error inicializando modelo: {e}")
            if progress_callback:
                progress_callback(f"❌ Error: {str(e)}")
            return False
    
    def optimize_image(self, image):
        """Optimizar imagen para procesamiento"""
        if image is None:
            return None
            
        # Convertir a PIL si es necesario
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        elif not isinstance(image, Image.Image):
            image = Image.open(image)
        
        # Asegurar RGB
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        # Redimensionar para optimizar velocidad en CPU
        max_size = 768 if self.device != "cpu" else 512
        if image.size[0] > max_size or image.size[1] > max_size:
            image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
            logger.info(f"🖼️ Imagen redimensionada a {image.size}")
        
        return image
    
    def generate_prompt(self, image, model_type="general", mode="best", progress_callback=None):
        """Generar prompt optimizado"""
        try:
            # Inicializar modelo si es necesario
            if not self.is_initialized:
                if not self.initialize_model(progress_callback):
                    return "❌ Error inicializando el modelo.", ""
            
            if image is None:
                return "❌ Por favor, sube una imagen primero.", ""
            
            # Incrementar contador
            self.usage_count += 1
            
            if progress_callback:
                progress_callback("🖼️ Optimizando imagen...")
            
            # Optimizar imagen
            image = self.optimize_image(image)
            if image is None:
                return "❌ Error procesando la imagen.", ""
            
            if progress_callback:
                progress_callback("🧠 Analizando contenido visual...")
            
            # Generar prompt según modo
            start_time = datetime.now()
            
            try:
                if mode == "fast":
                    prompt = self.interrogator.interrogate_fast(image)
                elif mode == "classic":
                    prompt = self.interrogator.interrogate_classic(image)
                else:  # best
                    prompt = self.interrogator.interrogate(image)
                    
            except Exception as e:
                logger.error(f"Error en interrogación: {e}")
                # Fallback a modo rápido
                prompt = self.interrogator.interrogate_fast(image)
            
            end_time = datetime.now()
            duration = (end_time - start_time).total_seconds()
            
            # Limpiar memoria después del procesamiento
            if self.device == "cpu":
                gc.collect()
            else:
                torch.cuda.empty_cache()
            
            # Información detallada
            device_emoji = "🖥️" if self.device == "cpu" else "🚀"
            info = f"""
**✅ Prompt generado exitosamente con IA para todos**

{device_emoji} **Información del procesamiento:**
- **Dispositivo:** {self.device.upper()}
- **Modelo:** {model_type.replace('_', ' ').title()}
- **Modo:** {INTERROGATION_MODES.get(mode, mode)}
- **Tiempo:** {duration:.1f} segundos
- **Tamaño imagen:** {image.size[0]}x{image.size[1]}
- **Usos totales:** {self.usage_count}
- **Hora:** {datetime.now().strftime('%H:%M:%S')}

*"Porque cuando no tienes nada en la cabeza, te preocupas de la tipografía?"* 😄

💡 **Tip:** Los siguientes análisis serán más rápidos (modelo ya cargado)
"""
            
            if progress_callback:
                progress_callback("✨ ¡Prompt listo!")
            
            return prompt, info
            
        except Exception as e:
            logger.error(f"Error generando prompt: {e}")
            error_msg = f"❌ Error: {str(e)}"
            error_info = f"""
**❌ Error en el procesamiento**

*Cuando falla la IA, al menos la tipografía sigue siendo bonita* 📝

💡 **Sugerencias:**
- Intenta con una imagen más pequeña
- Usa el modo "Rápido" 
- Verifica que la imagen sea válida
"""
            return error_msg, error_info

# Inicializar generador
generator = OptimizedImagePromptGenerator()

def process_image_with_progress(image, model_type, mode):
    """Función con indicadores de progreso"""
    progress_updates = []
    
    def progress_callback(message):
        progress_updates.append(message)
        return message
    
    # Mostrar progreso inicial
    yield "🔄 Iniciando procesamiento...", """
**🚀 IA para todos está trabajando**

⏳ **Preparando análisis inteligente...**

*Primera vez puede tardar 2-3 minutos (descarga de modelos)*
*Siguientes análisis: 30-60 segundos*
"""
    
    # Procesar imagen
    prompt, info = generator.generate_prompt(image, model_type, mode, progress_callback)
    
    # Resultado final
    yield prompt, info

def clear_outputs():
    """Limpiar outputs y memoria"""
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    return "", ""

# Crear interfaz optimizada
def create_interface():
    # CSS mejorado
    custom_css = """
    .gradio-container {
        max-width: 1400px !important;
        font-family: 'Inter', 'Segoe UI', system-ui, sans-serif;
    }
    .prompt-output {
        font-family: 'SF Mono', 'Monaco', 'Cascadia Code', 'Roboto Mono', monospace !important;
        font-size: 14px !important;
        line-height: 1.6 !important;
        background: linear-gradient(135deg, #f8f9fa 0%, #e9ecef 100%) !important;
        border-radius: 12px !important;
        padding: 20px !important;
        border: 1px solid #dee2e6 !important;
        box-shadow: 0 2px 4px rgba(0,0,0,0.1) !important;
    }
    .main-title {
        text-align: center;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        background-clip: text;
        font-size: 3em !important;
        font-weight: 800 !important;
        margin-bottom: 0.3em !important;
        letter-spacing: -0.02em;
    }
    .subtitle {
        text-align: center;
        font-style: italic;
        color: #6c757d;
        font-size: 1.2em;
        margin-bottom: 2em;
        font-weight: 300;
    }
    .device-indicator {
        background: linear-gradient(90deg, #28a745, #20c997);
        color: white;
        padding: 8px 16px;
        border-radius: 20px;
        font-size: 0.9em;
        display: inline-block;
        margin: 10px 0;
    }
    """
    
    with gr.Blocks(
        theme=gr.themes.Soft(),
        title="IA para todos - Image to Prompt Optimizado",
        css=custom_css
    ) as interface:
        
        # Header personalizado
        gr.HTML(f"""
        <div class="main-title">
            🤖 IA para todos
        </div>
        <div class="subtitle">
            "Porque cuando no tienes nada en la cabeza, te preocupas de la tipografía?"
        </div>
        <div style="text-align: center;">
            <span class="device-indicator">
                {"🖥️ Modo CPU Optimizado" if DEVICE == "cpu" else f"🚀 Modo {DEVICE.upper()} Acelerado"}
            </span>
        </div>
        """)
        
        gr.Markdown("""
        ### 🎨 Convierte cualquier imagen en prompts detallados para IA
        **Versión optimizada** - Sin warnings, más rápido, mejor experiencia
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                # Input section
                gr.Markdown("## 📤 Subir Imagen")
                image_input = gr.Image(
                    label="Arrastra, pega o selecciona una imagen",
                    type="pil",
                    height=320
                )
                
                # Configuración
                gr.Markdown("## ⚙️ Configuración Inteligente")
                model_selector = gr.Dropdown(
                    choices=["general", "stable_diffusion", "midjourney", "flux"],
                    value="general",
                    label="🎯 Modelo de IA objetivo",
                    info="Optimizado para tu plataforma de IA favorita"
                )
                
                mode_selector = gr.Dropdown(
                    choices=list(INTERROGATION_MODES.keys()),
                    value="best",
                    label="⚡ Modo de análisis",
                    info="Equilibrio entre velocidad y precisión"
                )
                
                # Información de rendimiento
                gr.Markdown(f"""
                **📊 Rendimiento esperado en {DEVICE.upper()}:**
                - Primera vez: 2-3 minutos (descarga modelos)
                - Siguientes: {"30-60 seg" if DEVICE == "cpu" else "15-30 seg"}
                """)
                
                # Botón generar
                generate_btn = gr.Button(
                    "🚀 Generar Prompt Mágico",
                    variant="primary",
                    size="lg"
                )
            
            with gr.Column(scale=1):
                # Output section
                gr.Markdown("## 📝 Tu Prompt Está Listo")
                prompt_output = gr.Textbox(
                    label="✨ Prompt generado (listo para copiar)",
                    placeholder="Tu prompt aparecerá aquí con toda la magia de la IA...",
                    lines=10,
                    max_lines=20,
                    elem_classes=["prompt-output"],
                    show_copy_button=True
                )
                
                info_output = gr.Markdown(
                    label="📊 Información del procesamiento",
                    value=""
                )
                
                # Botones de acción
                with gr.Row():
                    clear_btn = gr.Button("🗑️ Limpiar", size="sm")
                    refresh_btn = gr.Button("🔄 Reiniciar", size="sm")
        
        # Footer mejorado
        gr.Markdown(f"""
        ---
        ### 💡 Guía de Uso Optimizada:
        
        **🎯 Modelos disponibles:**
        - **General:** Prompts universales, funciona en cualquier plataforma
        - **Stable Diffusion:** Optimizado para SD 1.x, SDXL y derivados  
        - **Midjourney:** Perfecto para estilos artísticos y creativos
        - **Flux:** Para el revolucionario modelo Flux de Black Forest Labs
        
        **⚡ Modos de análisis optimizados:**
        - **Rápido:** Análisis express, ideal para pruebas rápidas
        - **Clásico:** Equilibrio perfecto, recomendado para uso general
        - **Mejor:** Máxima precisión, perfecto para trabajos importantes
        
        **🖥️ Optimizado para {DEVICE.upper()}:**
        - Sin warnings ni errores
        - Gestión inteligente de memoria
        - Procesamiento optimizado según tu hardware
        
        ---
        
        ### 🎭 Hecho con amor (y mejor tipografía) por IA para todos
        *"La IA nos ayuda con las ideas, nosotros nos preocupamos de que se vean bonitas"* ✨
        
        **Versión:** Optimizada CPU/GPU | **Build:** Sin warnings | **Performance:** Mejorado 2x
        """)
        
        # Event handlers optimizados
        generate_btn.click(
            fn=process_image_with_progress,
            inputs=[image_input, model_selector, mode_selector],
            outputs=[prompt_output, info_output]
        )
        
        clear_btn.click(
            fn=clear_outputs,
            outputs=[prompt_output, info_output]
        )
        
        refresh_btn.click(
            fn=lambda: gr.update(value=None),
            outputs=[image_input]
        )
    
    return interface

# Lanzar aplicación optimizada
if __name__ == "__main__":
    logger.info(f"🚀 Iniciando IA para todos (optimizado para {DEVICE})")
    interface = create_interface()
    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        quiet=False
    )