import os
import sys

# Activar la transferencia acelerada de Hugging Face antes de importar otros módulos
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"

import gradio as gr
import torch
from diffusers import DiffusionPipeline

# Intentar importar spaces para compatibilidad con Hugging Face Zero GPU
try:
    import spaces
    USING_SPACES = True
except ImportError:
    USING_SPACES = False

# Inicialización del modelo
# Al usar hf_transfer activo, la descarga del modelo de ~4.4GB se optimiza al máximo del ancho de banda disponible.
model_id = "prism-ml/bonsai-image-binary-4B-gemlite-1bit"

print("Cargando modelo... (Usando hf_transfer para máxima velocidad)")
pipe = DiffusionPipeline.from_pretrained(
    model_id, 
    dtype=torch.bfloat16, 
    device_map="auto" if not USING_SPACES else None
)

if not USING_SPACES and torch.cuda.is_available():
    pipe = pipe.to("cuda")

# CSS personalizado con estética Brutalista Moderna
brutalist_css = """
body, .gradio-container {
    background-color: #f3f3f3 !important;
    font-family: 'Courier New', Courier, monospace !important;
}
.brutal-card {
    border: 3px solid #000000 !important;
    box-shadow: 6px 6px 0px #000000 !important;
    background-color: #ffffff !important;
    border-radius: 0px !important;
    padding: 15px !important;
    margin-bottom: 15px !important;
}
.brutal-button {
    border: 3px solid #000000 !important;
    box-shadow: 4px 4px 0px #000000 !important;
    background-color: #ff5757 !important;
    color: #ffffff !important;
    font-weight: 900 !important;
    border-radius: 0px !important;
    text-transform: uppercase !important;
    letter-spacing: 1px !important;
    transition: all 0.1s ease !important;
}
.brutal-button:hover {
    background-color: #e04343 !important;
}
.brutal-button:active {
    box-shadow: 1px 1px 0px #000000 !important;
    transform: translate(3px, 3px) !important;
}
input, select, textarea, .secondary {
    border: 3px solid #000000 !important;
    border-radius: 0px !important;
    box-shadow: 3px 3px 0px #000000 !important;
    background-color: #ffffff !important;
    font-weight: bold !important;
}
.gr-form {
    border: none !important;
}
"""

def dynamic_style_inputs(style):
    """
    Controla la visibilidad de los paneles de opciones según el estilo seleccionado.
    """
    return (
        gr.update(visible=(style == "Realista")),
        gr.update(visible=(style == "Acuarela")),
        gr.update(visible=(style == "Pintura al Óleo")),
        gr.update(visible=(style == "Custom"))
    )

def calculate_dimensions(quality, aspect_ratio):
    """
    Calcula el ancho y alto basándose en la calidad y la relación de aspecto.
    """
    # Escala de píxeles base orientativa para evitar excesos de memoria
    base_pixels = 512 if quality == "0.5K" else (1024 if quality == "1K" else 1536) # 2K limitado a 1536 para estabilidad de VRAM
    
    if aspect_ratio == "1:1 (Cuadrado)":
        return base_pixels, base_pixels
    elif aspect_ratio == "16:9 (Horizontal)":
        return int(base_pixels * 1.25), int(base_pixels * 0.75)
    elif aspect_ratio == "9:16 (Vertical)":
        return int(base_pixels * 0.75), int(base_pixels * 1.25)
    return 1024, 1024

# Función contenedora para Zero GPU si aplica
def generate_inference(prompt, steps, width, height):
    if USING_SPACES:
        # En Zero GPU de Hugging Face se requiere decorar localmente o llamar dentro de un entorno compatible
        @spaces.GPU(duration=30)
        def run_on_space():
            pipe.to("cuda")
            return pipe(prompt, num_inference_steps=steps, width=width, height=height).images[0]
        return run_on_space()
    else:
        return pipe(prompt, num_inference_steps=steps, width=width, height=height).images[0]

def generate_image(
    prompt_base, style, 
    realista_light, realista_color, realista_texture, realista_comp,
    watercolor_paper, watercolor_wetness, watercolor_blend,
    oil_tech, oil_brush, oil_text,
    custom_text,
    quality, aspect_ratio, steps
):
    if not prompt_base.strip():
        return None, "Por favor, introduce un prompt base."
    
    # Construcción estructurada del prompt final según el estilo elegido
    full_prompt = prompt_base
    
    if style == "Realista":
        modifiers = f"{realista_light}, {realista_color}, {realista_texture} texture, {realista_comp} composition, highly realistic, photo"
        full_prompt = f"{full_prompt}, {modifiers}"
    elif style == "Acuarela":
        modifiers = f"watercolor painting, {watercolor_paper} paper, {watercolor_wetness} style, {watercolor_blend} colors, artistic"
        full_prompt = f"{full_prompt}, {modifiers}"
    elif style == "Pintura al Óleo":
        modifiers = f"oil painting, {oil_tech} technique, {oil_brush} brushstrokes, {oil_text} texture"
        full_prompt = f"{full_prompt}, {modifiers}"
    elif style == "Custom" and custom_text.strip():
        full_prompt = f"{full_prompt}, {custom_text}"
        
    width, height = calculate_dimensions(quality, aspect_ratio)
    
    # Asegurar dimensiones múltiplos de 8 para evitar errores en el backend de diffusers
    width = (width // 8) * 8
    height = (height // 8) * 8

    try:
        image = generate_inference(full_prompt, int(steps), width, height)
        return image, f"Prompt generado: {full_prompt}"
    except Exception as e:
        return None, f"Error durante la generación: {str(e)}"


# Construcción de la interfaz
with gr.Blocks(css=brutalist_css, title="Bonsai Studio") as demo:
    gr.HTML("<h1 style='text-align: center; font-weight: 900; margin-top: 20px; font-size: 2.5em;'>BONSAI STUDIO</h1>")
    gr.HTML("<p style='text-align: center; margin-bottom: 30px; font-weight: bold;'>Generación de Imágenes con Estilo Brutalista</p>")
    
    with gr.Row():
        with gr.Column(scale=1):
            with gr.Group(elem_classes="brutal-card"):
                prompt_input = gr.Textbox(
                    label="Prompt Base", 
                    placeholder="Ejemplo: Un astronauta en un bosque templado...", 
                    lines=3
                )
                
                style_select = gr.Dropdown(
                    choices=["Realista", "Acuarela", "Pintura al Óleo", "Custom"], 
                    value="Realista", 
                    label="Estilo Artístico"
                )
            
            # Sub-opciones de estilo dinámicas
            with gr.Group(elem_classes="brutal-card"):
                # Opciones para Realista
                with gr.Column(visible=True) as col_realista:
                    gr.Markdown("### Opciones de Realismo")
                    r_light = gr.Dropdown(choices=["Estudio", "Natural", "Cinematográfica", "Luz de Neón", "Contraluz"], value="Cinematográfica", label="Iluminación")
                    r_color = gr.Dropdown(choices=["Vibrante", "Cálido", "Frío", "Desaturado", "Monocromo"], value="Vibrante", label="Paleta de Colores")
                    r_texture = gr.Dropdown(choices=["Porosa/Micro-detalles", "Piel Suave", "Imperfecto/Orgánico", "Pulido/Liso"], value="Porosa/Micro-detalles", label="Texturas")
                    r_comp = gr.Dropdown(choices=["Plano Detalle", "Primer Plano", "Plano Medio", "Gran Plano General"], value="Plano Medio", label="Composición")

                # Opciones para Acuarela
                with gr.Column(visible=False) as col_watercolor:
                    gr.Markdown("### Opciones de Acuarela")
                    w_paper = gr.Dropdown(choices=["Prensado en Frío", "Texturizado Rugoso", "Papel Suave Japones"], value="Prensado en Frío", label="Tipo de Papel")
                    w_wetness = gr.Dropdown(choices=["Mojado sobre Mojado (Líquido)", "Pincel Seco (Detalle)", "Aguada Uniforme"], value="Mojado sobre Mojado (Líquido)", label="Efecto de Humedad")
                    w_blend = gr.Dropdown(choices=["Sangrado Suave", "Bordes Definidos", "Gradiente de Color"], value="Sangrado Suave", label="Mezcla de Colores")

                # Opciones para Pintura al Óleo
                with gr.Column(visible=False) as col_oil:
                    gr.Markdown("### Opciones de Pintura al Óleo")
                    o_tech = gr.Dropdown(choices=["Impasto (Espeso)", "Glaseado (Capas finas)", "Alla Prima (Directo)"], value="Impasto (Espeso)", label="Técnica de Pintura")
                    o_brush = gr.Dropdown(choices=["Pinceladas Gruesas y Visibles", "Pinceladas Difuminadas", "Espátula de Pintor"], value="Pinceladas Gruesas y Visibles", label="Patrón de Pinceladas")
                    o_text = gr.Dropdown(choices=["Muy Texturizado", "Relieve Medio", "Superficie Plana"], value="Muy Texturizado", label="Textura Física")

                # Opción para Custom
                with gr.Column(visible=False) as col_custom:
                    gr.Markdown("### Modificadores Manuales")
                    custom_text = gr.Textbox(placeholder="Añade modificadores personalizados separados por comas...", label="Modificadores de Prompt", lines=2)

            with gr.Group(elem_classes="brutal-card"):
                gr.Markdown("### Configuración Técnica")
                quality_select = gr.Dropdown(choices=["0.5K", "1K", "2K"], value="1K", label="Calidad/Resolución")
                aspect_ratio_select = gr.Dropdown(choices=["1:1 (Cuadrado)", "16:9 (Horizontal)", "9:16 (Vertical)"], value="1:1 (Cuadrado)", label="Relación de Aspecto")
                steps_slider = gr.Slider(minimum=1, maximum=8, value=4, step=1, label="Pasos de Inferencia (Steps)")
                
            generate_btn = gr.Button("GENERAR IMAGEN", elem_classes="brutal-button")

        with gr.Column(scale=1):
            with gr.Group(elem_classes="brutal-card"):
                output_image = gr.Image(label="Imagen Generada", type="pil")
                output_status = gr.Textbox(label="Información del Prompt Aplicado", interactive=False)

    # Conectar lógica de cambio de visibilidad de componentes
    style_select.change(
        fn=dynamic_style_inputs,
        inputs=[style_select],
        outputs=[col_realista, col_watercolor, col_oil, col_custom]
    )

    # Acción del botón generar
    generate_btn.click(
        fn=generate_image,
        inputs=[
            prompt_input, style_select,
            r_light, r_color, r_texture, r_comp,
            w_paper, w_wetness, w_blend,
            o_tech, o_brush, o_text,
            custom_text,
            quality_select, aspect_ratio_select, steps_slider
        ],
        outputs=[output_image, output_status]
    )

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