import gradio as gr
from core.generation_logic import run_cn_preprocessor_entry
from core.pipelines.controlnet_preprocessor import CPU_ONLY_PREPROCESSORS
from utils.app_utils import PREPROCESSOR_MODEL_MAP, PREPROCESSOR_PARAMETER_MAP
from ui.layout import MAX_DYNAMIC_CONTROLS

def attach_event_handlers(ui_components, demo):
    def update_cn_input_visibility(choice):
        return {
            ui_components["cn_image_input"]: gr.update(visible=choice == "Image"),
            ui_components["cn_video_input"]: gr.update(visible=choice == "Video")
        }
    ui_components["cn_input_type"].change(
        fn=update_cn_input_visibility,
        inputs=[ui_components["cn_input_type"]],
        outputs=[ui_components["cn_image_input"], ui_components["cn_video_input"]]
    )
    
    def update_preprocessor_models_dropdown(preprocessor_name):
        models = PREPROCESSOR_MODEL_MAP.get(preprocessor_name)
        if models:
            model_filenames = [m[1] for m in models]
            return gr.update(choices=model_filenames, value=model_filenames[0], visible=True)
        else:
            return gr.update(choices=[], value=None, visible=False)

    def update_preprocessor_settings_ui(preprocessor_name):
        params = PREPROCESSOR_PARAMETER_MAP.get(preprocessor_name, [])
        
        slider_updates, dropdown_updates, checkbox_updates = [], [], []
        
        s_idx, d_idx, c_idx = 0, 0, 0

        for param in params:
            if s_idx + d_idx + c_idx >= MAX_DYNAMIC_CONTROLS: break
            
            name = param["name"]
            ptype = param["type"]
            config = param["config"]
            label = name.replace('_', ' ').title()

            if ptype == "INT" or ptype == "FLOAT":
                if s_idx < MAX_DYNAMIC_CONTROLS:
                    slider_updates.append(gr.update(
                        label=label,
                        minimum=config.get('min', 0),
                        maximum=config.get('max', 255),
                        step=config.get('step', 0.1 if ptype == "FLOAT" else 1),
                        value=config.get('default', 0),
                        visible=True
                    ))
                    s_idx += 1
            elif isinstance(ptype, list):
                if d_idx < MAX_DYNAMIC_CONTROLS:
                    dropdown_updates.append(gr.update(
                        label=label,
                        choices=ptype,
                        value=config.get('default', ptype[0] if ptype else None),
                        visible=True
                    ))
                    d_idx += 1
            elif ptype == "BOOLEAN":
                if c_idx < MAX_DYNAMIC_CONTROLS:
                    checkbox_updates.append(gr.update(
                        label=label,
                        value=config.get('default', False),
                        visible=True
                    ))
                    c_idx += 1

        for _ in range(s_idx, MAX_DYNAMIC_CONTROLS): slider_updates.append(gr.update(visible=False))
        for _ in range(d_idx, MAX_DYNAMIC_CONTROLS): dropdown_updates.append(gr.update(visible=False))
        for _ in range(c_idx, MAX_DYNAMIC_CONTROLS): checkbox_updates.append(gr.update(visible=False))

        return slider_updates + dropdown_updates + checkbox_updates

    def update_run_button_for_cpu(preprocessor_name):
        if preprocessor_name in CPU_ONLY_PREPROCESSORS:
            return gr.update(value="Run Preprocessor CPU Only", variant="primary"), gr.update(visible=False)
        else:
            return gr.update(value="Run Preprocessor", variant="primary"), gr.update(visible=True)

    ui_components["preprocessor_cn"].change(
        fn=update_preprocessor_models_dropdown,
        inputs=[ui_components["preprocessor_cn"]],
        outputs=[ui_components["preprocessor_model_cn"]]
    ).then(
        fn=update_preprocessor_settings_ui,
        inputs=[ui_components["preprocessor_cn"]],
        outputs=ui_components["cn_sliders"] + ui_components["cn_dropdowns"] + ui_components["cn_checkboxes"]
    ).then(
        fn=update_run_button_for_cpu,
        inputs=[ui_components["preprocessor_cn"]],
        outputs=[ui_components["run_cn"], ui_components["zero_gpu_cn"]]
    )

    all_dynamic_inputs = (
        ui_components["cn_sliders"] + 
        ui_components["cn_dropdowns"] + 
        ui_components["cn_checkboxes"]
    )

    ui_components["run_cn"].click(
        fn=run_cn_preprocessor_entry,
        inputs=[
            ui_components["cn_input_type"],
            ui_components["cn_image_input"],
            ui_components["cn_video_input"],
            ui_components["preprocessor_cn"],
            ui_components["preprocessor_model_cn"],
            ui_components["zero_gpu_cn"],
        ] + all_dynamic_inputs,
        outputs=[ui_components["output_gallery_cn"]]
    )

    def run_on_load():
        all_updates = {}
        
        default_preprocessor = "Canny Edge" 
        model_update = update_preprocessor_models_dropdown(default_preprocessor)
        all_updates[ui_components["preprocessor_model_cn"]] = model_update
        
        settings_outputs = update_preprocessor_settings_ui(default_preprocessor)
        dynamic_outputs = ui_components["cn_sliders"] + ui_components["cn_dropdowns"] + ui_components["cn_checkboxes"]
        for i, comp in enumerate(dynamic_outputs):
            all_updates[comp] = settings_outputs[i]

        run_button_update, zero_gpu_update = update_run_button_for_cpu(default_preprocessor)
        all_updates[ui_components["run_cn"]] = run_button_update
        all_updates[ui_components["zero_gpu_cn"]] = zero_gpu_update
        
        return all_updates
    
    all_load_outputs = [
        ui_components["preprocessor_model_cn"],
        *ui_components["cn_sliders"],
        *ui_components["cn_dropdowns"],
        *ui_components["cn_checkboxes"],
        ui_components["run_cn"],
        ui_components["zero_gpu_cn"]
    ]

    demo.load(
        fn=run_on_load,
        outputs=all_load_outputs
    )