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 build_ui(event_handler_function):
    ui_components = {}

    # Helper function to load thumbnails from the output folder.
    # It returns a **deduplicated** list of absolute image file paths.
    # Using a set ensures that even if the same filename appears multiple
    # times (e.g., due to accidental duplicate saves), the Gallery tab will
    # only display each image once.
    def load_gallery_images():
        """Return a deduplicated list of image paths.
        If multiple files share the same stem (e.g., ``img.png`` and ``img.webp``),
        only the first encountered file (based on sorted order) is kept. This prevents
        the Gallery tab from displaying duplicate visual entries that arise from
        different extensions of the same generated image.
        """
        output_dir = os.path.abspath(OUTPUT_DIR)
        if not os.path.isdir(output_dir):
            return []
        seen_bases = set()
        image_paths = []
        for fname in sorted(os.listdir(output_dir)):
            low = fname.lower()
            if low.endswith((".png", ".jpg", ".jpeg", ".gif", ".webp")):
                base = os.path.splitext(fname)[0]
                if base not in seen_bases:
                    seen_bases.add(base)
                    image_paths.append(os.path.join(output_dir, fname))
        return image_paths

    # All UI components, including tabs, must be created inside the same Gradio Blocks context.
    with gr.Blocks() as demo:
        gr.Markdown("# ImageGen")
        gr.Markdown(
            "This demo is a streamlined version of the [Comfy web UI](https://github.com/RioShiina47/comfy-webui)'s [ImageGen](https://huggingface.co/spaces/RioShiina/ImageGen) functionality. "
            "Other spaces: [ImageGen1](https://huggingface.co/spaces/RioShiina/ImageGen1), "
            "[ImageGen2](https://huggingface.co/spaces/RioShiina/ImageGen2), "
            "[ImageGen3](https://huggingface.co/spaces/RioShiina/ImageGen3), "
            "[ImageGen4](https://huggingface.co/spaces/RioShiina/ImageGen4), "
            "[ImageGen5](https://huggingface.co/spaces/RioShiina/ImageGen5), "
            "[ImageGen6](https://huggingface.co/spaces/RioShiina/ImageGen6), "
            "[ImageGen7](https://huggingface.co/spaces/RioShiina/ImageGen7), "
            "[ImageGen8](https://huggingface.co/spaces/RioShiina/ImageGen8)"
        )

        # Tabs container – now correctly nested within the Blocks context.
        with gr.Tabs(elem_id="tabs_container") as 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())

            # New Gallery tab to display generated images
            with gr.TabItem("Gallery", id=5):
                # Refresh button to manually reload thumbnails from the output folder.
                refresh_btn = gr.Button("Refresh Gallery", variant="secondary")
                gallery = gr.Gallery(label="Generated Images", show_label=False, columns=3, height="auto", type="filepath")
                ui_components["gallery"] = gallery
                # Bind the refresh button to the same loader used on startup.
                refresh_btn.click(fn=load_gallery_images, outputs=gallery)

            ui_components["tabs"] = tabs
            ui_components["image_gen_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>")

        # Load gallery images on startup using the helper defined earlier.
        demo.load(fn=load_gallery_images, inputs=[], outputs=ui_components["gallery"])

        # Connect event handlers (e.g., button clicks) with the UI components.
        event_handler_function(ui_components, demo)

    return demo