import gradio as gr
from core.settings import MODEL_MAP_CHECKPOINT
from .ui_components import (
    create_base_parameter_ui, 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
)

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))
        
        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")
        
        components[f'prompt_{prefix}'] = gr.Text(label="Prompt", lines=3)
        components[f'neg_prompt_{prefix}'] = gr.Text(label="Negative prompt", lines=3)
        
        with gr.Row():
            with gr.Column(scale=1):
                param_defaults = {'w': 1024, 'h': 1024, 'cs_vis': False, 'cs_val': 1}
                components.update(create_base_parameter_ui(prefix, param_defaults))
            with gr.Column(scale=1):
                components[f'result_{prefix}'] = gr.Gallery(label="Result", show_label=False, columns=2, object_fit="contain", 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))
        
    return components