import gradio as gr
from core.settings import MODEL_MAP_CHECKPOINT, RESOLUTION_MAP
from comfy_integration.nodes import SAMPLER_CHOICES, SCHEDULER_CHOICES
from .ui_components import (
    create_lora_settings_ui,
    create_controlnet_ui, create_anima_controlnet_lllite_ui, create_diffsynth_controlnet_ui, create_ipadapter_ui, create_embedding_ui,
    create_conditioning_ui, create_vae_override_ui,
    create_model_architecture_filter_ui, create_category_filter_ui,
    create_sd3_ipadapter_ui, create_flux1_ipadapter_ui, create_style_ui,
    create_reference_latent_ui, create_hidream_o1_reference_ui,
    create_pid_ui,
    DEFAULT_SAMPLER, DEFAULT_SCHEDULER, DEFAULT_STEPS, DEFAULT_CFG,
)

def create_ui():
    """Creates the UI components for the Txt2Img tab."""
    prefix = "txt2img"
    components = {}
    
    with gr.Column():
        components.update(create_model_architecture_filter_ui(prefix))
        
        # Layout for model selection and run button
        with gr.Row():
            components.update(create_category_filter_ui(prefix))
            components[f'base_model_{prefix}'] = gr.Dropdown(
                label="Base Model",
                choices=list(MODEL_MAP_CHECKPOINT.keys()),
                value=list(MODEL_MAP_CHECKPOINT.keys())[0],
                scale=3,
                allow_custom_value=True,
            )
            with gr.Column(scale=1):
                components[f'run_{prefix}'] = gr.Button("Run", variant="primary")
        
        # Prompt fields
        components[f'prompt_{prefix}'] = gr.Text(label="Prompt", lines=3)
        components[f'neg_prompt_{prefix}'] = gr.Text(label="Negative prompt", lines=3)
        
        # Layout: left column – parameter blocks; right column – result gallery (same height as left)
        with gr.Row():
            # LEFT: all parameter blocks stacked vertically
            with gr.Column():
                # Aspect Ratio block
                components[f'aspect_ratio_{prefix}'] = gr.Dropdown(
                    label="Aspect Ratio",
                    choices=list(RESOLUTION_MAP.get('sdxl', {}).keys()),
                    value="1:1 (Square)",
                    interactive=True,
                    allow_custom_value=True,
                )
                # Width & Height block
                with gr.Row():
                    components[f'width_{prefix}'] = gr.Number(label="Width", value=1024, interactive=True)
                    components[f'height_{prefix}'] = gr.Number(label="Height", value=1024, interactive=True)
                # Sampler & Scheduler block
                with gr.Row():
                    components[f'sampler_{prefix}'] = gr.Dropdown(
                        label="Sampler",
                        choices=SAMPLER_CHOICES,
                        value=DEFAULT_SAMPLER if DEFAULT_SAMPLER in SAMPLER_CHOICES else (SAMPLER_CHOICES[0] if SAMPLER_CHOICES else 'euler')
                    )
                    components[f'scheduler_{prefix}'] = gr.Dropdown(
                        label="Scheduler",
                        choices=SCHEDULER_CHOICES,
                        value=DEFAULT_SCHEDULER if DEFAULT_SCHEDULER in SCHEDULER_CHOICES else (SCHEDULER_CHOICES[0] if SCHEDULER_CHOICES else 'simple')
                    )
                # Steps & CFG Scale block
                with gr.Row():
                    components[f'steps_{prefix}'] = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=DEFAULT_STEPS)
                    components[f'cfg_{prefix}'] = gr.Slider(label="CFG Scale", minimum=1.0, maximum=20.0, step=0.1, value=DEFAULT_CFG)
                # Seed & Batch Size block
                with gr.Row():
                    components[f'seed_{prefix}'] = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
                    components[f'batch_size_{prefix}'] = gr.Slider(label="Batch Size", minimum=1, maximum=16, step=1, value=1)
                # ZeroGPU Duration block
                components[f'zero_gpu_{prefix}'] = gr.Number(
                    label="ZeroGPU Duration (s)",
                    value=None,
                    placeholder="Default: 60s, Max: 120s",
                    info="Optional: Set how long to reserve the GPU."
                )
            # RIGHT: result gallery – height matches combined left column height
            with gr.Column():
                components[f'result_{prefix}'] = gr.Gallery(
                    label="Result",
                    show_label=False,
                    columns=2,
                    type="filepath",
                    height=627
                )
        

        components.update(create_lora_settings_ui(prefix))
        components.update(create_controlnet_ui(prefix))
        components.update(create_anima_controlnet_lllite_ui(prefix))
        components.update(create_diffsynth_controlnet_ui(prefix))
        components.update(create_ipadapter_ui(prefix))
        components.update(create_flux1_ipadapter_ui(prefix))
        components.update(create_sd3_ipadapter_ui(prefix))
        components.update(create_embedding_ui(prefix))
        components.update(create_style_ui(prefix))
        components.update(create_conditioning_ui(prefix))
        components.update(create_reference_latent_ui(prefix))
        components.update(create_hidream_o1_reference_ui(prefix))
        components.update(create_vae_override_ui(prefix))
        components.update(create_pid_ui(prefix))
        
    return components