import gradio as gr
import torch
from diffusers import DiffusionPipeline
import gc 


print("💼 Iniciando BizSketch. Cargando motor gráfico ligero...")
model_id = "SimianLuo/LCM_Dreamshaper_v7"

try:

    pipe = DiffusionPipeline.from_pretrained(model_id)
    pipe.to("cpu", torch.float32)
    pipe.safety_checker = None 
    print("✅ Motor listo. Modo de bajo consumo activado.")

except Exception as e:
    print(f"❌ Error crítico al cargar el modelo: {e}\nAsegúrate de tener instalado: pip install diffusers transformers accelerate torch")
    pipe = None

def limpiar_memoria():
    gc.collect()

# --- 2. LÓGICA DE GENERACIÓN ---
def generar_concepto(concepto_negocio, estilo_visual):
    if pipe is None:
        raise gr.Error("El modelo no se pudo cargar. Revisa la terminal.")
    
    if not concepto_negocio:
         raise gr.Error("Por favor, escribe un concepto primero.")

    prompt_base = f"{concepto_negocio}, {estilo_visual}, professional corporate illustration, clean background, minimal design, high quality, presentation ready"
    
    prompt_negativo = "blurry, messy, amateur, ugly, deformed hands, text, watermark, signature, low resolution, photograph"

    print(f"📊 Generando concepto: {concepto_negocio}...")
    
    try:

        imagen = pipe(
            prompt=prompt_base,
            negative_prompt=prompt_negativo,
            num_inference_steps=5,  
            guidance_scale=7.5,
            width=512,                               
            height=512
        ).images[0]
        
        limpiar_memoria() 
        return imagen
    except Exception as e:
        limpiar_memoria()
        raise gr.Error(f"Error durante la generación: {e}")

# --- 3. DISEÑO DE LA INTERFAZ (Layout Profesional) ---

tema_negocio = gr.themes.Monochrome(
    primary_hue="slate",
    radius_size=gr.themes.sizes.radius_sm,
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
)

css_corporativo = """
.container {max-width: 900px; margin: auto; padding-top: 20px;}
#header-box {text-align: center; margin-bottom: 30px; border-bottom: 2px solid #e5e7eb; padding-bottom: 20px;}
h1 {font-weight: 700; color: #1a202c;}
.subtitle {color: #4b5563; font-size: 1.1rem;}
#generate-btn {font-weight: bold; font-size: 1.1rem;}
"""

with gr.Blocks(theme=tema_negocio, css=css_corporativo, title="BizSketch App") as demo:
    
    with gr.Column(elem_classes="container"):
        # Header
        with gr.Column(elem_id="header-box"):
            gr.Markdown("# 💼 BizSketch: Visualizador de Conceptos")
            gr.Markdown("<span class='subtitle'>Genera ilustraciones rápidas para tus presentaciones y brainstorms.</span>")

        with gr.Row(equal_height=False):
            # --- Columna Izquierda: Inputs ---
            with gr.Column(scale=2):
                with gr.Group():
                    gr.Markdown("### 1. Define tu idea")
                    input_concept = gr.Textbox(
                        label="Concepto de Negocio",
                        placeholder="Ej: Crecimiento sostenible en mercados emergentes...",
                        lines=3,
                        show_label=False
                    )
                
                with gr.Group():
                     gr.Markdown("### 2. Elige el estilo visual")
                     radio_style = gr.Radio(
                        choices=[
                            "clean line art sketch, blueprint style", 
                            "flat vector illustration, modern tech style", 
                            "isometric 3D render, soft colors",
                            "watercolor sketch, business storyboard style"
                        ],
                        value="flat vector illustration, modern tech style",
                        label="Estilo",
                        show_label=False,
                        interactive=True
                    )

                btn = gr.Button("🚀 Generar Visualización", variant="primary", elem_id="generate-btn")
                
                with gr.Accordion("💡 Consejos para mejores resultados", open=False):
                    gr.Markdown("- Mantén los conceptos simples (ej: 'reunión de equipo exitosa' en vez de una descripción de 3 párrafos).")
                    gr.Markdown("- El estilo 'Flat vector' suele ser el más limpio para presentaciones.")

            # --- Columna Derecha: Output ---
            with gr.Column(scale=3):
                 gr.Markdown("### 🎯 Resultado")
                 output_image = gr.Image(
                    label="Visualización Generada",
                    type="pil",
                    interactive=False,
                    #show_download_button=True,
                    height=400
                )

        # --- Footer / Ejemplos ---
        gr.Markdown("---")
        gr.Markdown("### ⚡ Ejemplos Rápidos")
        gr.Examples(
            examples=[
                ["Equipo diverso colaborando alrededor de una mesa con gráficos", "flat vector illustration, modern tech style"],
                ["Innovación tecnológica y medio ambiente, bombilla con hoja", "clean line art sketch, blueprint style"],
                ["Logística global, red de conexiones alrededor del mundo", "isometric 3D render, soft colors"]
            ],
            inputs=[input_concept, radio_style],
            outputs=output_image,
            fn=generar_concepto,
            cache_examples=False 
        )

    btn.click(fn=generar_concepto, inputs=[input_concept, radio_style], outputs=output_image)

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