import gradio as gr
from PIL import Image
from super_image import EdsrModel, ImageLoader
import os
import uuid

# Cargamos el modelo profesional (EDSR Base x4)
model = EdsrModel.from_pretrained('eugenesiow/edsr-base', scale=4)

def enhance_image(input_img):
    if input_img is None:
        return None
    
    # 1. Asegurar que la entrada sea una imagen PIL
    if not isinstance(input_img, Image.Image):
        input_img = Image.fromarray(input_img)
    
    # 2. Convertir la imagen para la IA
    inputs = ImageLoader.load_image(input_img)
    
    # 3. Procesar la imagen con la IA
    preds = model(inputs)
    
    # 4. LA SOLUCIÓN AL CUADRO NEGRO:
    # Usamos el método nativo de la librería para guardar la imagen.
    # Generamos un nombre único para evitar conflictos si varias personas lo usan a la vez.
    unique_filename = f"resultado_{uuid.uuid4().hex}.png"
    
    # Esta función interna sabe exactamente cómo convertir los números de la IA a colores reales
    ImageLoader.save_image(preds, unique_filename)
    
    # Devolvemos la ruta del archivo para que Gradio lo muestre
    return unique_filename

# Interfaz Premium
with gr.Blocks() as demo:
    gr.HTML("<h1 style='text-align:center;'>📸 AI Photo Enhancer Ultra</h1>")
    gr.Markdown("<p style='text-align:center;'>Procesamiento x4 de alta definición.</p>")
    
    with gr.Row():
        with gr.Column():
            # Entrada de imagen
            input_i = gr.Image(label="Imagen Original", type="pil", sources=["upload"])
            btn = gr.Button("✨ MEJORAR CALIDAD", variant="primary")
        
        with gr.Column():
            # Salida: Ahora espera la ruta de un archivo (filepath) en lugar de datos puros
            output_i = gr.Image(label="Resultado HD", type="filepath", interactive=False)

    # Conexión del botón
    btn.click(
        fn=enhance_image, 
        inputs=input_i, 
        outputs=output_i
    )

if __name__ == "__main__":
    demo.launch(ssr_mode=False)