import gradio as gr
import torch
from PIL import Image
import spaces
from diffusers import QwenImageEditPipeline
import os
import traceback
from datetime import datetime


device = "cuda" if torch.cuda.is_available() else "cpu"

model_repo_id = "Qwen/Qwen-Image-Edit"

if torch.cuda.is_available():
    torch_dtype = torch.bfloat16
else:
    torch_dtype = torch.float32

def log_msg(mensaje):
    """Logger con timestamp"""
    timestamp = datetime.now().strftime("%H:%M:%S")
    print(f"[{timestamp}] {mensaje}")
    return mensaje

log_msg("🔄 Iniciando carga del modelo Qwen-Image-Edit...")

pipe = QwenImageEditPipeline.from_pretrained(
    model_repo_id, 
    torch_dtype=torch_dtype
)

log_msg("✅ Modelo cargado correctamente en CPU/RAM")

@spaces.GPU(duration=90)
def editar(
    imagen_path,
    prompt,
    seed,
    randomize_seed,
    guidance_scale,
    num_inference_steps,
    progress=gr.Progress(track_tqdm=True),
):
    """Función de edición con LOGGING COMPLETO"""
    
    try:
        # ============ PASO 1: VALIDACIÓN DE ENTRADA ============
        log_msg("=" * 60)
        log_msg("🚀 INICIANDO NUEVA EDICIÓN")
        log_msg("=" * 60)
        
        if imagen_path is None:
            error_msg = "❌ No se proporcionó imagen"
            log_msg(error_msg)
            raise gr.Error(error_msg)
        
        log_msg(f"✅ Imagen recibida: {imagen_path}")
        
        # Verificar que el archivo existe
        if not os.path.exists(imagen_path):
            error_msg = f"❌ Archivo no encontrado: {imagen_path}"
            log_msg(error_msg)
            raise gr.Error(error_msg)
        
        log_msg(f"✅ Archivo existe (tamaño: {os.path.getsize(imagen_path) / 1024:.2f} KB)")
        
        # ============ PASO 2: CARGAR Y VALIDAR IMAGEN ============
        log_msg("📸 Cargando imagen...")
        
        try:
            imagen_original = Image.open(imagen_path)
            log_msg(f"✅ Imagen cargada: {imagen_original.size} - Modo: {imagen_original.mode}")
        except Exception as e:
            error_msg = f"❌ Error al abrir imagen: {str(e)}"
            log_msg(error_msg)
            raise gr.Error(error_msg)
        
        # Convertir a RGB si es necesario
        if imagen_original.mode != "RGB":
            log_msg(f"🔄 Convirtiendo de {imagen_original.mode} a RGB...")
            imagen_original = imagen_original.convert("RGB")
            log_msg("✅ Conversión completada")
        
        # ============ PASO 3: REDIMENSIONAR SI ES MUY GRANDE ============
        width, height = imagen_original.size
        max_dimension = 1024  # Límite para evitar OOM
        
        if width > max_dimension or height > max_dimension:
            log_msg(f"⚠️ Imagen muy grande ({width}x{height}), redimensionando...")
            
            # Calcular nuevo tamaño manteniendo aspect ratio
            if width > height:
                new_width = max_dimension
                new_height = int((height / width) * max_dimension)
            else:
                new_height = max_dimension
                new_width = int((width / height) * max_dimension)
            
            imagen_original = imagen_original.resize(
                (new_width, new_height), 
                Image.Resampling.LANCZOS
            )
            log_msg(f"✅ Redimensionada a {new_width}x{new_height}")
        else:
            log_msg(f"✅ Tamaño OK: {width}x{height}")
        
        # ============ PASO 4: VALIDAR PROMPT ============
        if not prompt or prompt.strip() == "":
            error_msg = "❌ Prompt vacío"
            log_msg(error_msg)
            raise gr.Error("Por favor describe qué quieres editar")
        
        log_msg(f"✅ Prompt: '{prompt}'")
        
        # ============ PASO 5: CONFIGURAR SEMILLA ============
        if randomize_seed:
            import random
            seed = random.randint(0, 2147483647)
        
        log_msg(f"🎲 Semilla: {seed}")
        
        generator = torch.Generator().manual_seed(seed)
        
        # ============ PASO 6: MOVER MODELO A GPU ============
        log_msg("🔄 Moviendo modelo a GPU...")
        
        try:
            pipe.to(device)
            log_msg(f"✅ Modelo en {device}")
        except Exception as e:
            error_msg = f"❌ Error al mover a GPU: {str(e)}"
            log_msg(error_msg)
            raise gr.Error(error_msg)
        
        # Verificar memoria GPU
        if torch.cuda.is_available():
            memoria_libre = torch.cuda.mem_get_info()[0] / 1024**3
            memoria_total = torch.cuda.mem_get_info()[1] / 1024**3
            log_msg(f"💾 GPU Memoria: {memoria_libre:.2f}GB libre / {memoria_total:.2f}GB total")
        
        # ============ PASO 7: EJECUTAR PIPELINE ============
        log_msg(f"🎨 Iniciando edición (pasos={num_inference_steps}, cfg={guidance_scale})...")
        
        try:
            resultado = pipe(
                image=imagen_original,
                prompt=prompt,
                generator=generator,
                true_cfg_scale=guidance_scale,
                num_inference_steps=num_inference_steps,
                negative_prompt=""
            )
            log_msg("✅ Pipeline completado exitosamente")
            
        except torch.cuda.OutOfMemoryError:
            error_msg = "❌ OUT OF MEMORY - Imagen muy grande o pasos muy altos"
            log_msg(error_msg)
            torch.cuda.empty_cache()
            raise gr.Error("💾 Error: GPU sin memoria. Intenta con una imagen más pequeña o menos pasos.")
            
        except Exception as e:
            error_msg = f"❌ Error en pipeline: {str(e)}"
            log_msg(error_msg)
            log_msg(f"Traceback completo:\n{traceback.format_exc()}")
            raise gr.Error(f"Error durante edición: {str(e)}")
        
        # ============ PASO 8: PROCESAR RESULTADO ============
        if resultado.images and len(resultado.images) > 0:
            imagen_final = resultado.images[0]
            log_msg(f"✅ Imagen resultado generada: {imagen_final.size}")
            log_msg("=" * 60)
            log_msg("🎉 EDICIÓN COMPLETADA EXITOSAMENTE")
            log_msg("=" * 60)
            
            return imagen_final, seed
        else:
            error_msg = "❌ No se generó ninguna imagen"
            log_msg(error_msg)
            raise gr.Error("Error: el modelo no generó resultado")
    
    except gr.Error:
        raise  # Re-lanzar errores de Gradio
    
    except Exception as e:
        error_msg = f"❌ ERROR INESPERADO: {str(e)}"
        log_msg(error_msg)
        log_msg(f"Traceback:\n{traceback.format_exc()}")
        raise gr.Error(f"Error inesperado: {str(e)}")
    
    finally:
        # Limpiar cache GPU
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            log_msg("🧹 Cache GPU limpiado")


examples = [
    "Change the car color to red",
    "Add text 'SALE' in large red letters",
    "Change the background to a beach sunset",
    "Make it look like a watercolor painting",
    "Remove the background",
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 960px;
}
"""

with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 🎨 Editor de Imágenes con Qwen-Image-Edit (DEBUG MODE)
            ### 🔍 Versión con logging completo para debugging
            
            **Revisa la consola/logs para ver el progreso detallado**
            """
        )

        with gr.Row():
            with gr.Column(scale=1):
                imagen_input = gr.Image(
                    label="📸 Imagen Original",
                    type="filepath",
                    height=400,
                    sources=["upload", "clipboard"]
                )
                
                prompt = gr.Textbox(
                    label="✍️ Describe la edición",
                    show_label=True,
                    max_lines=3,
                    placeholder="Ejemplo: Change the color to blue",
                    lines=3
                )

                run_button = gr.Button("🚀 Editar Imagen", variant="primary", size="lg")

            with gr.Column(scale=1):
                result = gr.Image(
                    label="✨ Resultado", 
                    show_label=True, 
                    height=400
                )

        with gr.Accordion("⚙️ Configuración", open=True):
            
            with gr.Row():
                num_inference_steps = gr.Slider(
                    label="Pasos (⚠️ menos pasos = más rápido, menos OOM)",
                    minimum=20,
                    maximum=50,  # ⚠️ REDUCIDO para evitar timeout
                    step=5,
                    value=30,  # ⚠️ DEFAULT más bajo
                )
                
                guidance_scale = gr.Slider(
                    label="CFG Scale",
                    minimum=1.0,
                    maximum=10.0,
                    step=0.5,
                    value=3.5,
                )
            
            with gr.Row():
                seed = gr.Slider(
                    label="Semilla",
                    minimum=0,
                    maximum=2147483647,
                    step=1,
                    value=42,
                )

                randomize_seed = gr.Checkbox(label="Semilla aleatoria", value=True)

        gr.Markdown("### 💡 Ejemplos:")
        gr.Examples(
            examples=examples, 
            inputs=[prompt]
        )
        
        gr.Markdown(
            """
            ---
            ### 🔍 DEBUG INFO:
            - Todos los pasos se registran en los logs
            - Las imágenes grandes se redimensionan automáticamente (max 1024px)
            - Se limpie el cache GPU después de cada uso
            
            ### ⚠️ Si ves "ZeroGPU worker error":
            1. **Imagen muy grande** → Se auto-redimensiona, pero revisa logs
            2. **Timeout >90s** → Reduce pasos de inferencia (usa 30 en vez de 50)
            3. **OOM** → Imagen muy pesada o muchos pasos
            
            ### 📊 Qué revisar en los logs:
            - ✅ Todos los pasos deben mostrar checkmarks verdes
            - 💾 Memoria GPU disponible
            - 📸 Tamaño de imagen procesada
            - ⏱️ Si tarda mucho, reduce pasos
            """
        )
    
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=editar,
        inputs=[
            imagen_input,
            prompt,
            seed,
            randomize_seed,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, seed],
    )

if __name__ == "__main__":
    demo.queue(max_size=20)
    demo.launch()