import gradio as gr
from .chain_handlers import (
    create_lora_event_handlers,
    create_controlnet_event_handlers,
    create_anima_controlnet_lllite_event_handlers,
    create_diffsynth_controlnet_event_handlers,
    create_ipadapter_event_handlers,
    create_embedding_event_handlers,
    create_conditioning_event_handlers,
    create_flux1_ipadapter_event_handlers,
    create_sd3_ipadapter_event_handlers,
    create_style_event_handlers,
    create_reference_latent_event_handlers,
    create_hidream_o1_reference_event_handlers
)
from .change_handlers import (
    make_update_fn,
    make_model_change_fn,
    run_on_load,
    on_aspect_ratio_change
)
from .run_handlers import create_run_event

def attach_event_handlers(ui_components, demo):
    for prefix, task_type in [
        ("txt2img", "txt2img"), ("img2img", "img2img"), ("inpaint", "inpaint"),
        ("outpaint", "outpaint"), ("hires_fix", "hires_fix"),
    ]:

        arch_comp = ui_components.get(f'model_arch_{prefix}')
        cat_comp = ui_components.get(f'model_cat_{prefix}')
        model_comp = ui_components.get(f'base_model_{prefix}')
        clip_skip_comp = ui_components.get(f'clip_skip_{prefix}') or ui_components.get(f'{prefix}_clip_skip')
        guidance_comp = ui_components.get(f'guidance_{prefix}') or ui_components.get(f'{prefix}_guidance')
        aspect_ratio_comp = ui_components.get(f'aspect_ratio_{prefix}') or ui_components.get(f'{prefix}_aspect_ratio_dropdown')
        width_comp = ui_components.get(f'width_{prefix}') or ui_components.get(f'{prefix}_width')
        height_comp = ui_components.get(f'height_{prefix}') or ui_components.get(f'{prefix}_height')
        
        cn_types_list = ui_components.get(f'controlnet_types_{prefix}', [])
        cn_series_list = ui_components.get(f'controlnet_series_{prefix}', [])
        cn_filepaths_list = ui_components.get(f'controlnet_filepaths_{prefix}', [])
        
        anima_cn_types_list = ui_components.get(f'anima_controlnet_lllite_types_{prefix}', [])
        anima_cn_series_list = ui_components.get(f'anima_controlnet_lllite_series_{prefix}', [])
        anima_cn_filepaths_list = ui_components.get(f'anima_controlnet_lllite_filepaths_{prefix}', [])

        diffsynth_cn_types_list = ui_components.get(f'diffsynth_controlnet_types_{prefix}', [])
        diffsynth_cn_series_list = ui_components.get(f'diffsynth_controlnet_series_{prefix}', [])
        diffsynth_cn_filepaths_list = ui_components.get(f'diffsynth_controlnet_filepaths_{prefix}', [])
        
        lora_accordion = ui_components.get(f'lora_accordion_{prefix}')
        cn_accordion = ui_components.get(f'controlnet_accordion_{prefix}')
        anima_cn_accordion = ui_components.get(f'anima_controlnet_lllite_accordion_{prefix}')
        diffsynth_cn_accordion = ui_components.get(f'diffsynth_controlnet_accordion_{prefix}')
        ipa_accordion = ui_components.get(f'ipadapter_accordion_{prefix}')
        sd3_ipa_accordion = ui_components.get(f'sd3_ipadapter_accordion_{prefix}')
        flux1_ipa_accordion = ui_components.get(f'flux1_ipadapter_accordion_{prefix}')
        style_accordion = ui_components.get(f'style_accordion_{prefix}')
        embedding_accordion = ui_components.get(f'embedding_accordion_{prefix}')
        conditioning_accordion = ui_components.get(f'conditioning_accordion_{prefix}')
        ref_latent_accordion = ui_components.get(f'reference_latent_accordion_{prefix}')
        hidream_o1_ref_accordion = ui_components.get(f'hidream_o1_reference_accordion_{prefix}')
        pid_accordion = ui_components.get(f'pid_accordion_{prefix}')
        vae_accordion = ui_components.get(f'vae_accordion_{prefix}')
        
        ipa_preset_list = ui_components.get(f'ipadapter_final_preset_{prefix}')
        
        prompt_comp = ui_components.get(f'prompt_{prefix}') or ui_components.get(f'{prefix}_positive_prompt')
        neg_prompt_comp = ui_components.get(f'neg_prompt_{prefix}') or ui_components.get(f'{prefix}_negative_prompt')
        steps_comp = ui_components.get(f'steps_{prefix}') or ui_components.get(f'{prefix}_steps')
        cfg_comp = ui_components.get(f'cfg_{prefix}') or ui_components.get(f'{prefix}_cfg')
        sampler_comp = ui_components.get(f'sampler_{prefix}') or ui_components.get(f'{prefix}_sampler_name')
        scheduler_comp = ui_components.get(f'scheduler_{prefix}') or ui_components.get(f'{prefix}_scheduler')
        
        extra_comps = [prompt_comp, neg_prompt_comp, steps_comp, cfg_comp, sampler_comp, scheduler_comp, width_comp, height_comp]
        valid_extra_comps = [c for c in extra_comps if c is not None]
        
        if arch_comp and cat_comp and model_comp:
            outputs = [model_comp, cat_comp]
            if clip_skip_comp: outputs.append(clip_skip_comp)
            if guidance_comp: outputs.append(guidance_comp)
            if aspect_ratio_comp: outputs.append(aspect_ratio_comp)
            outputs.extend(cn_types_list + cn_series_list + cn_filepaths_list)
            outputs.extend(anima_cn_types_list + anima_cn_series_list + anima_cn_filepaths_list)
            outputs.extend(diffsynth_cn_types_list + diffsynth_cn_series_list + diffsynth_cn_filepaths_list)
            if lora_accordion: outputs.append(lora_accordion)
            if cn_accordion: outputs.append(cn_accordion)
            if anima_cn_accordion: outputs.append(anima_cn_accordion)
            if diffsynth_cn_accordion: outputs.append(diffsynth_cn_accordion)
            if ipa_accordion: outputs.append(ipa_accordion)
            if sd3_ipa_accordion: outputs.append(sd3_ipa_accordion)
            if flux1_ipa_accordion: outputs.append(flux1_ipa_accordion)
            if style_accordion: outputs.append(style_accordion)
            if embedding_accordion: outputs.append(embedding_accordion)
            if conditioning_accordion: outputs.append(conditioning_accordion)
            if ref_latent_accordion: outputs.append(ref_latent_accordion)
            if hidream_o1_ref_accordion: outputs.append(hidream_o1_ref_accordion)
            if pid_accordion: outputs.append(pid_accordion)
            if vae_accordion: outputs.append(vae_accordion)
            if ipa_preset_list: outputs.append(ipa_preset_list)
            
            outputs.extend(valid_extra_comps)
            
            update_fn = make_update_fn(
                model_comp, cat_comp, clip_skip_comp, aspect_ratio_comp, width_comp, height_comp, 
                cn_types_list, cn_series_list, cn_filepaths_list, 
                anima_cn_types_list, anima_cn_series_list, anima_cn_filepaths_list,
                diffsynth_cn_types_list, diffsynth_cn_series_list, diffsynth_cn_filepaths_list, 
                ipa_preset_list, lora_accordion, cn_accordion, anima_cn_accordion, diffsynth_cn_accordion, ipa_accordion, sd3_ipa_accordion, flux1_ipa_accordion, style_accordion, embedding_accordion, conditioning_accordion,
                ref_latent_accordion, hidream_o1_ref_accordion, guidance_comp, prompt_comp, neg_prompt_comp, steps_comp, cfg_comp, sampler_comp, scheduler_comp,
                pid_acc=pid_accordion, vae_acc=vae_accordion
            )
            inputs = [arch_comp, cat_comp]
            if aspect_ratio_comp:
                inputs.append(aspect_ratio_comp)
            arch_comp.change(fn=update_fn, inputs=inputs, outputs=outputs)
            cat_comp.change(fn=update_fn, inputs=inputs, outputs=outputs)

        if model_comp:
            outputs2 = []
            if arch_comp: outputs2.append(arch_comp)
            if cat_comp: outputs2.append(cat_comp)
            if clip_skip_comp: outputs2.append(clip_skip_comp)
            if guidance_comp: outputs2.append(guidance_comp)
            if aspect_ratio_comp: outputs2.append(aspect_ratio_comp)
            outputs2.extend(cn_types_list + cn_series_list + cn_filepaths_list)
            outputs2.extend(anima_cn_types_list + anima_cn_series_list + anima_cn_filepaths_list)
            outputs2.extend(diffsynth_cn_types_list + diffsynth_cn_series_list + diffsynth_cn_filepaths_list)
            if lora_accordion: outputs2.append(lora_accordion)
            if cn_accordion: outputs2.append(cn_accordion)
            if anima_cn_accordion: outputs2.append(anima_cn_accordion)
            if diffsynth_cn_accordion: outputs2.append(diffsynth_cn_accordion)
            if ipa_accordion: outputs2.append(ipa_accordion)
            if sd3_ipa_accordion: outputs2.append(sd3_ipa_accordion)
            if flux1_ipa_accordion: outputs2.append(flux1_ipa_accordion)
            if style_accordion: outputs2.append(style_accordion)
            if embedding_accordion: outputs2.append(embedding_accordion)
            if conditioning_accordion: outputs2.append(conditioning_accordion)
            if ref_latent_accordion: outputs2.append(ref_latent_accordion)
            if hidream_o1_ref_accordion: outputs2.append(hidream_o1_ref_accordion)
            if pid_accordion: outputs2.append(pid_accordion)
            if vae_accordion: outputs2.append(vae_accordion)
            if ipa_preset_list: outputs2.append(ipa_preset_list)
            
            outputs2.extend(valid_extra_comps)
            
            if outputs2:
                inputs2 = [model_comp]
                if arch_comp: inputs2.append(arch_comp)
                if cat_comp: inputs2.append(cat_comp)
                if aspect_ratio_comp: inputs2.append(aspect_ratio_comp)
                change_fn = make_model_change_fn(
                    cat_comp, clip_skip_comp, aspect_ratio_comp, width_comp, height_comp, 
                    cn_types_list, cn_series_list, cn_filepaths_list, 
                    anima_cn_types_list, anima_cn_series_list, anima_cn_filepaths_list,
                    diffsynth_cn_types_list, diffsynth_cn_series_list, diffsynth_cn_filepaths_list, 
                    arch_comp, ipa_preset_list, lora_accordion, cn_accordion, anima_cn_accordion, diffsynth_cn_accordion, ipa_accordion, sd3_ipa_accordion, flux1_ipa_accordion, style_accordion, embedding_accordion, conditioning_accordion,
                    ref_latent_accordion, hidream_o1_ref_accordion, guidance_comp, prompt_comp, neg_prompt_comp, steps_comp, cfg_comp, sampler_comp, scheduler_comp,
                    pid_acc=pid_accordion, vae_acc=vae_accordion
                )
                model_comp.change(fn=change_fn, inputs=inputs2, outputs=outputs2)

        create_lora_event_handlers(prefix, ui_components)
        create_controlnet_event_handlers(prefix, ui_components)
        create_anima_controlnet_lllite_event_handlers(prefix, ui_components)
        create_diffsynth_controlnet_event_handlers(prefix, ui_components)
        create_ipadapter_event_handlers(prefix, ui_components)
        create_embedding_event_handlers(prefix, ui_components)
        create_conditioning_event_handlers(prefix, ui_components)
        create_flux1_ipadapter_event_handlers(prefix, ui_components)
        create_sd3_ipadapter_event_handlers(prefix, ui_components)
        create_style_event_handlers(prefix, ui_components)
        create_reference_latent_event_handlers(prefix, ui_components)
        create_hidream_o1_reference_event_handlers(prefix, ui_components)
        create_run_event(prefix, task_type, ui_components)

    if 'view_mode_inpaint' in ui_components:
        def toggle_inpaint_fullscreen_view(view_mode):
            is_fullscreen = (view_mode == "Fullscreen View")
            other_elements_visible = not is_fullscreen
            editor_height = 800 if is_fullscreen else 272
            
            updates = {
                ui_components['prompts_column_inpaint']: gr.update(visible=other_elements_visible),
                ui_components['params_and_gallery_row_inpaint']: gr.update(visible=other_elements_visible),
                ui_components['accordion_wrapper_inpaint']: gr.update(visible=other_elements_visible),
                ui_components['input_image_dict_inpaint']: gr.update(height=editor_height),
            }
            
            model_and_run_rows = ui_components.get('model_and_run_row_inpaint', [])
            for row in model_and_run_rows:
                updates[row] = gr.update(visible=other_elements_visible)
            
            return updates
        
        output_components = []
        model_and_run_rows = ui_components.get('model_and_run_row_inpaint', [])
        if isinstance(model_and_run_rows, list):
            output_components.extend(model_and_run_rows)
        else:
            output_components.append(model_and_run_rows)
        
        output_components.extend([
            ui_components['prompts_column_inpaint'],
            ui_components['params_and_gallery_row_inpaint'], 
            ui_components['accordion_wrapper_inpaint'],
            ui_components['input_image_dict_inpaint']
        ])
        
        ui_components['view_mode_inpaint'].change(
            fn=toggle_inpaint_fullscreen_view, 
            inputs=[ui_components['view_mode_inpaint']], 
            outputs=output_components, 
            show_progress=False
        )

    all_load_outputs = []
    for prefix in ["txt2img", "img2img", "inpaint", "outpaint", "hires_fix"]:
        if f'controlnet_types_{prefix}' in ui_components:
            all_load_outputs.extend(ui_components[f'controlnet_types_{prefix}'])
            all_load_outputs.extend(ui_components[f'controlnet_series_{prefix}'])
            all_load_outputs.extend(ui_components[f'controlnet_filepaths_{prefix}'])
        if f'anima_controlnet_lllite_types_{prefix}' in ui_components:
            all_load_outputs.extend(ui_components[f'anima_controlnet_lllite_types_{prefix}'])
            all_load_outputs.extend(ui_components[f'anima_controlnet_lllite_series_{prefix}'])
            all_load_outputs.extend(ui_components[f'anima_controlnet_lllite_filepaths_{prefix}'])
        if f'diffsynth_controlnet_types_{prefix}' in ui_components:
            all_load_outputs.extend(ui_components[f'diffsynth_controlnet_types_{prefix}'])
            all_load_outputs.extend(ui_components[f'diffsynth_controlnet_series_{prefix}'])
            all_load_outputs.extend(ui_components[f'diffsynth_controlnet_filepaths_{prefix}'])
        if f'ipadapter_final_preset_{prefix}' in ui_components:
            all_load_outputs.extend(ui_components[f'ipadapter_lora_strengths_{prefix}'])
            all_load_outputs.append(ui_components[f'ipadapter_final_preset_{prefix}'])
            all_load_outputs.append(ui_components[f'ipadapter_final_lora_strength_{prefix}'])

    if all_load_outputs:
        demo.load(
            fn=lambda: run_on_load(ui_components),
            outputs=all_load_outputs
        )

    for prefix in ["txt2img", "img2img", "inpaint", "outpaint", "hires_fix"]:
        aspect_ratio_dropdown = ui_components.get(f'aspect_ratio_{prefix}') or ui_components.get(f'{prefix}_aspect_ratio_dropdown')
        width_component = ui_components.get(f'width_{prefix}') or ui_components.get(f'{prefix}_width')
        height_component = ui_components.get(f'height_{prefix}') or ui_components.get(f'{prefix}_height')
        model_dropdown = ui_components.get(f'base_model_{prefix}')
        if aspect_ratio_dropdown and width_component and height_component and model_dropdown:
            aspect_ratio_dropdown.change(fn=on_aspect_ratio_change, inputs=[aspect_ratio_dropdown, model_dropdown], outputs=[width_component, height_component], show_progress=False)