import os
import gradio as gr
from core.settings import *

from .shared import txt2img_ui, img2img_ui, inpaint_ui, outpaint_ui, hires_fix_ui

MAX_DYNAMIC_CONTROLS = 10

def get_preprocessor_choices():
    from nodes import NODE_DISPLAY_NAME_MAPPINGS
    
    preprocessor_names = [
        display_name for class_name, display_name in NODE_DISPLAY_NAME_MAPPINGS.items()
        if "Preprocessor" in class_name or "Segmentor" in class_name or 
           "Estimator" in class_name or "Detector" in class_name
    ]
    return sorted(list(set(preprocessor_names)))


def build_ui(event_handler_function):
    ui_components = {}

    with gr.Blocks() as demo:
        gr.Markdown("# ImageGen - Z-Image")
        gr.Markdown(
            "This demo is a streamlined version of the [Comfy web UI](https://github.com/RioShiina47/comfy-webui)'s ImgGen functionality. "
            "Other versions are also available: "
            "[Qwen-Image](https://huggingface.co/spaces/RioShiina/ImageGen-Qwen-Image), "
            "[Illstrious](https://huggingface.co/spaces/RioShiina/ImageGen-Illstrious), "
            "[NoobAI](https://huggingface.co/spaces/RioShiina/ImageGen-NoobAI), "
            "[Pony](https://huggingface.co/spaces/RioShiina/ImageGen-Pony1), "
            "[SDXL](https://huggingface.co/spaces/RioShiina/ImageGen-SDXL), "
            "[SD1.5](https://huggingface.co/spaces/RioShiina/ImageGen-SD15)"
        )
        with gr.Tabs(elem_id="tabs_container") as tabs:
            with gr.TabItem("Z-Image", id=0):
                with gr.Tabs(elem_id="image_gen_tabs") as image_gen_tabs:
                    with gr.TabItem("Txt2Img", id=0):
                        ui_components.update(txt2img_ui.create_ui())
                    
                    with gr.TabItem("Img2Img", id=1):
                        ui_components.update(img2img_ui.create_ui())

                    with gr.TabItem("Inpaint", id=2):
                        ui_components.update(inpaint_ui.create_ui())

                    with gr.TabItem("Outpaint", id=3):
                        ui_components.update(outpaint_ui.create_ui())

                    with gr.TabItem("Hires. Fix", id=4):
                        ui_components.update(hires_fix_ui.create_ui())
                
                ui_components['image_gen_tabs'] = image_gen_tabs

            with gr.TabItem("Controlnet Preprocessors", id=1):
                gr.Markdown("## ControlNet Auxiliary Preprocessors")
                gr.Markdown("Powered by [Fannovel16/comfyui_controlnet_aux](https://github.com/Fannovel16/comfyui_controlnet_aux).")
                gr.Markdown("Upload an image or video to process it with a ControlNet preprocessor.")
                with gr.Row():
                    with gr.Column(scale=1):
                        cn_input_type = gr.Radio(["Image", "Video"], label="Input Type", value="Image")
                        cn_image_input = gr.Image(type="pil", label="Input Image", visible=True, height=384)
                        cn_video_input = gr.Video(label="Input Video", visible=False)
                        preprocessor_cn = gr.Dropdown(label="Preprocessor", choices=get_preprocessor_choices(), value="Canny Edge")
                        preprocessor_model_cn = gr.Dropdown(label="Preprocessor Model", choices=[], value=None, visible=False)
                        with gr.Column() as preprocessor_settings_ui:
                            cn_sliders, cn_dropdowns, cn_checkboxes = [], [], []
                            for i in range(MAX_DYNAMIC_CONTROLS):
                                cn_sliders.append(gr.Slider(visible=False, label=f"dyn_slider_{i}"))
                                cn_dropdowns.append(gr.Dropdown(visible=False, label=f"dyn_dropdown_{i}"))
                                cn_checkboxes.append(gr.Checkbox(visible=False, label=f"dyn_checkbox_{i}"))
                        run_cn = gr.Button("Run Preprocessor", variant="primary")
                    with gr.Column(scale=1):
                        output_gallery_cn = gr.Gallery(label="Output", show_label=False, object_fit="contain", height=512)
                        zero_gpu_cn = gr.Number(label="ZeroGPU Duration (s)", value=None, placeholder="Default: 60s, Max: 120s", info="Optional")
                ui_components.update({
                    "cn_input_type": cn_input_type, "cn_image_input": cn_image_input, "cn_video_input": cn_video_input,
                    "preprocessor_cn": preprocessor_cn, "preprocessor_model_cn": preprocessor_model_cn, "run_cn": run_cn,
                    "zero_gpu_cn": zero_gpu_cn, "output_gallery_cn": output_gallery_cn,
                    "preprocessor_settings_ui": preprocessor_settings_ui, "cn_sliders": cn_sliders,
                    "cn_dropdowns": cn_dropdowns, "cn_checkboxes": cn_checkboxes
                })
            
            ui_components["tabs"] = tabs
        
        gr.Markdown("<div style='text-align: center; margin-top: 20px;'>Made by RioShiina with ❤️<br><a href='https://github.com/RioShiina47' target='_blank'>GitHub</a> | <a href='https://huggingface.co/RioShiina' target='_blank'>Hugging Face</a> | <a href='https://civitai.com/user/RioShiina' target='_blank'>Civitai</a></div>")
        
        event_handler_function(ui_components, demo)

    return demo