import spaces
import gradio as gr
from PIL import Image
import io
import os

# Importa a função melhorada do seu módulo (que já está ajustado )
from faceless_processing import build_faceless_embroidery_assets, LineartConfig

@spaces.GPU
def process_faceless(input_image, prompt="", negative_prompt="", is_cartoon=False, detail="balanced"):
    if input_image is None:
        return None, None

    # Configuração base
    cfg = LineartConfig()
    
    # The detail preset applies to portraits too. Previously it only affected
    # cartoon mode, which made the public controls misleading.
    if detail == "detailed":
        cfg.lineart_threshold = 72
        cfg.lineart_min_perimeter = 20.0
        cfg.lineart_min_area = 5.0
        cfg.min_line_distance = 4
        cfg.silhouette_line_thickness = 5
        cfg.interior_line_thickness = 4
    elif detail == "soft":
        cfg.lineart_threshold = 100
        cfg.lineart_min_perimeter = 36.0
        cfg.lineart_min_area = 10.0
        cfg.min_line_distance = 7
        cfg.silhouette_line_thickness = 4
        cfg.interior_line_thickness = 3
    else:
        cfg.lineart_threshold = 88
        cfg.lineart_min_perimeter = 28.0
        cfg.lineart_min_area = 8.0
        cfg.min_line_distance = 5
        cfg.silhouette_line_thickness = 5
        cfg.interior_line_thickness = 4

    if is_cartoon:
        cfg.is_cartoon = True
        # Afrouxa as regras para não apagar detalhes do desenho, mas com limites seguros para bordar
        cfg.lineart_threshold = 30           # Deixa os traços mais fracos passarem
        cfg.lineart_min_perimeter = 25.0      # Evita traços minúsculos inferiores a 25px
        cfg.lineart_min_area = 10.0
        cfg.lineart_morph_open = False
        cfg.lineart_coarse = False           # OBRIGA a IA a enxergar detalhes finos
        cfg.force_canny = True               # Força a usar o detector de bordas
        cfg.min_contour_perimeter = 20.0      # Evita ruídos insignificantes de Canny, mantendo olhos/pupilas
        
        # Filtro de proximidade entre linhas para o modo cartoon manual
        cfg.min_line_distance = 6
        cfg.line_overlap_threshold = 0.40
        cfg.custom_tuned_for_embroidery = True
        
        # Ajuste dinâmico com base no nível de detalhe (Detail Level)
        if detail == "detailed":
            cfg.gaussian_ksize = 3
            cfg.bilateral_d = 3
            cfg.canny_low = 30
            cfg.canny_high = 70
            cfg.min_line_distance = 4        # Um pouco mais denso no modo detailed
        elif detail == "soft":
            cfg.gaussian_ksize = 7
            cfg.bilateral_d = 9
            cfg.canny_low = 60
            cfg.canny_high = 120
            cfg.min_line_distance = 8        # Mais espaçado no modo soft
        else: # balanced
            cfg.gaussian_ksize = 5
            cfg.bilateral_d = 5
            cfg.canny_low = 45
            cfg.canny_high = 95
            cfg.min_line_distance = 6
        
        cfg.detect_eyebrows = False
        cfg.erase_faces = False
        cfg.erase_borders = False
        
        # Impede de achar rosto e apagar os olhos/boca do desenho
        cfg.face_min_size = (9999, 9999)
        
        # Mantém grossinho para bordar!
        cfg.stroke_dilate_iterations = 0

    # NÃO REDIMENSIONAMOS A IMAGEM – a função interna ( embroidery_processing.py) já lida com tamanhos variados. 
    # Gera lineart e cutout usando a função otimizada
    try:
        lineart, cutout = build_faceless_embroidery_assets(
            input_image,
            cfg=cfg,
            use_huggingface=False,   # Desativa Hugging Face (pode ativar se tiver token)
            prompt=prompt or None,
            negative_prompt=negative_prompt or None,
        )
        return lineart, cutout
    except Exception as e:
        print(f"[ERRO] process_faceless: {e}")
        # Fallback: retorna a imagem original e uma mensagem de erro
        return input_image, None

# Interface Gradio
demo = gr.Interface(
    fn=process_faceless,
    inputs=[
        gr.Image(type="pil", label="Imagem de Entrada"),
        gr.Textbox(label="Prompt (Opcional)", value=""),
        gr.Textbox(label="Prompt Negativo (Opcional)", value=""),
        gr.Checkbox(label="Modo Cartoon / Desenho (Mantém detalhes e rostos)", value=False),
        gr.Dropdown(choices=["soft", "balanced", "detailed"], value="balanced", label="Nível de Detalhe (Apenas Modo Cartoon)")
    ],
    outputs=[
        gr.Image(label="Lineart Final (Transparente)"),
        gr.Image(label="Recorte do Rosto (Sem Fundo)")
    ],
    title="Faceless Embroidery Studio Generator",
    description="Gera contornos e linearts minimalistas (Faceless) para bordados a partir de retratos.\n\n**Nota:** O rosto é completamente apagado e apenas a silhueta externa é desenhada (exceto no Modo Cartoon)."
)

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