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_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
)

def create_ui():
    prefix = "inpaint"
    components = {}
    
    with gr.Column():
        with gr.Row() as arch_row:
            components.update(create_model_architecture_filter_ui(prefix))

        with gr.Row() as model_and_run_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
            )
            with gr.Column(scale=1): 
                components[f'run_{prefix}'] = gr.Button("Run Inpaint", variant="primary")
        
        components[f'model_and_run_row_{prefix}'] = [arch_row, model_and_run_row]

        with gr.Row() as main_content_row:
            with gr.Column(scale=1) as editor_column:
                components[f'view_mode_{prefix}'] = gr.Radio(
                    ["Normal View", "Fullscreen View"], 
                    label="Editor View", 
                    value="Normal View", 
                    interactive=True
                )
                components[f'input_image_dict_{prefix}'] = gr.ImageEditor(
                    type="pil", 
                    label="Image & Mask", 
                    height=272
                )
            components[f'editor_column_{prefix}'] = editor_column

            with gr.Column(scale=2) as prompts_column:
                components[f'prompt_{prefix}'] = gr.Text(label="Prompt", lines=6)
                components[f'neg_prompt_{prefix}'] = gr.Text(label="Negative prompt", lines=6)
            components[f'prompts_column_{prefix}'] = prompts_column

        with gr.Row() as params_and_gallery_row:
            with gr.Column(scale=1):
                from comfy_integration.nodes import SAMPLER_CHOICES, SCHEDULER_CHOICES
                with gr.Row():
                    components[f'denoise_{prefix}'] = gr.Slider(
                        label="Denoise", minimum=0.0, maximum=1.0, step=0.05, value=1.0
                    )
                    components[f'grow_mask_by_{prefix}'] = gr.Slider(
                        label="Grow Mask By", minimum=0, maximum=64, step=1, value=6
                    )
                with gr.Row():
                    components[f'sampler_{prefix}'] = gr.Dropdown(label="Sampler", choices=SAMPLER_CHOICES, value=SAMPLER_CHOICES[0])
                    components[f'scheduler_{prefix}'] = gr.Dropdown(label="Scheduler", choices=SCHEDULER_CHOICES, value='normal' if 'normal' in SCHEDULER_CHOICES else SCHEDULER_CHOICES[0])
                with gr.Row():
                    components[f'steps_{prefix}'] = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=28)
                    components[f'cfg_{prefix}'] = gr.Slider(label="CFG Scale", minimum=1.0, maximum=20.0, step=0.1, value=7.5)
                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)
                with gr.Row():
                    components[f'clip_skip_{prefix}'] = gr.Slider(label="Clip Skip", minimum=1, maximum=2, step=1, value=1, visible=False, interactive=True)
                    components[f'guidance_{prefix}'] = gr.Slider(label="Guidance (FLUX)", minimum=1.0, maximum=10.0, step=0.1, value=3.5, visible=False, interactive=True)
                    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.")

                components[f'width_{prefix}'] = gr.State(value=512)
                components[f'height_{prefix}'] = gr.State(value=512)

            with gr.Column(scale=1):
                components[f'result_{prefix}'] = gr.Gallery(label="Result", show_label=False, columns=1, object_fit="contain", height=510)
        
        components[f'params_and_gallery_row_{prefix}'] = params_and_gallery_row
            
        with gr.Column() as accordion_wrapper:

            components.update(create_lora_settings_ui(prefix))
            components.update(create_controlnet_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_style_ui(prefix))
            components.update(create_embedding_ui(prefix))
            components.update(create_conditioning_ui(prefix))
            components.update(create_reference_latent_ui(prefix))
            components.update(create_vae_override_ui(prefix))
        components[f'accordion_wrapper_{prefix}'] = accordion_wrapper
            
    return components