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_ipadapter_ui, create_embedding_ui,
    create_conditioning_ui, create_vae_override_ui, create_api_key_ui
)

def create_ui():
    """Creates the UI components for the Txt2Img tab."""
    prefix = "txt2img"
    components = {}
    
    with gr.Column():
        with gr.Row():
            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)
            with gr.Column(scale=1): 
                components[f'run_{prefix}'] = gr.Button("Run", variant="primary")
        
        components[f'prompt_{prefix}'] = gr.Text(label="Prompt", lines=3, placeholder="Enter your prompt")
        components[f'neg_prompt_{prefix}'] = gr.Text(label="Negative prompt", lines=3, value="monochrome, (low quality, worst quality:1.2), 3d, watermark, signature, ugly, poorly drawn,")
        
        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_api_key_ui(prefix))
        components.update(create_lora_settings_ui(prefix))
        components.update(create_controlnet_ui(prefix))
        components.update(create_ipadapter_ui(prefix))
        components.update(create_embedding_ui(prefix))
        components.update(create_conditioning_ui(prefix))
        components.update(create_vae_override_ui(prefix))
        
    return components