# app.py

# --- 🚨 CORREÇÃO DE EMERGÊNCIA: Garantir websockets >= 11.0.3 ANTES de importar qualquer coisa ---
import subprocess
import sys

def install_package(package):
    subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-cache-dir", package])

try:
    import websockets
    # Verifica se o submódulo 'asyncio' existe (presente só a partir da v11+)
    if not hasattr(websockets, 'asyncio'):
        raise ImportError("Versão antiga do websockets detectada. Atualizando...")
except (ImportError, AttributeError):
    print("📦 Instalando websockets>=11.0.3...")
    install_package("websockets>=11.0.3")
    # Recarrega o módulo após instalação
    import importlib
    import websockets
    importlib.reload(websockets)
    print("✅ websockets atualizado com sucesso!")
# --- ✅ FIM DA CORREÇÃO ---

# Agora sim, importe o resto
import os
import gradio as gr
from utils.editor import run_pipeline
from PIL import Image
import io

DESCRIPTION = """
# 👗 Fashion Editor — Stable Diffusion + ControlNet (OpenPose)
Envie a foto da modelo e uma foto da peça (jaqueta, camiseta etc.). A ferramenta
irá alinhar a peça, extrair a pose e usar Stable Diffusion + ControlNet para gerar
a composição fotorrealista.
"""

examples = [
    # Você pode adicionar exemplos aqui (caminhos relativos dentro do repo) após fazer upload
]

def infer(model_image, garment_image, prompt_extra):
    """
    model_image: PIL.Image or numpy
    garment_image: PIL.Image or numpy
    prompt_extra: str - description tweaks
    """
    if model_image is None or garment_image is None:
        return None, "Envie as duas imagens (modelo e peça)."

    # run_pipeline retorna (pil_image_result, debug_info_dict)
    try:
        result, info = run_pipeline(model_image, garment_image, prompt_extra)
        return result, f"OK — steps: {info.get('steps')}, guidance: {info.get('guidance_scale')}"
    except Exception as e:
        return None, f"Erro durante a geração: {str(e)}"

with gr.Blocks(title="Fashion Editor (SD + ControlNet)") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column():
            model_input = gr.Image(label="Foto da Modelo", type="pil")
            garment_input = gr.Image(label="Foto da Peça (Roupa)", type="pil")
            prompt_extra = gr.Textbox(label="Prompt extra (opcional) — ex: 'leather biker jacket, black, shiny, realistic'", lines=2)
            run_button = gr.Button("Gerar Edição")
            status = gr.Textbox(label="Status / Informações", interactive=False)
        with gr.Column():
            output = gr.Image(label="Resultado (Final)")

    run_button.click(fn=infer, inputs=[model_input, garment_input, prompt_extra], outputs=[output, status])

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