import gradio as gr
from core.model_loader import generate_caption

def build_ui(model, processor):

    def run_caption(image, style):
        if image is None:
            return "Please upload an image."
        image = image.convert("RGB")
        return generate_caption(
            model=model,
            processor=processor,
            image=image,
            style=style
        )

    with gr.Blocks(
        title="BLIP Image Captioning",
        fill_height=True
    ) as demo:

        gr.Markdown("## 🖼️ BLIP Image Captioning")
        gr.Markdown(
            "Generate captions or explanations from images using a CPU-only BLIP model (Salesforce/blip-image-captioning-large)."
        )

        # -----------------------------
        # MAIN FIXED LAYOUT
        # -----------------------------
        with gr.Row(equal_height=True):

            # LEFT: Image Column (Fixed)
            with gr.Column(scale=1, min_width=420):
                image_input = gr.Image(
                    type="pil",
                    label="Upload Image",
                    height=420,
                    elem_id="image-box"
                )
                generate_btn = gr.Button(
                    "Generate",
                    variant="primary"
                )

            # RIGHT: Controls + Output (Fixed)
            with gr.Column(scale=1, min_width=420, elem_id="right-panel"):
                style_select = gr.Dropdown(
                    choices=[
                        "Short Caption",
                        "Detailed Caption"
                    ],
                    value="Detailed Caption",
                    label="Caption Style"
                )

                output_text = gr.Textbox(
                    label="Generated Output",
                    lines=6,
                    max_lines=8,
                    elem_id="output-box"
                )

        # -----------------------------
        # EXAMPLES (SEPARATE SECTION)
        # -----------------------------
        gr.Markdown("### Examples")

        gr.Examples(
            examples=[
                ["./assets/zebra.jpg", "Short Caption"],
                ["./assets/cat.jpg", "Short Caption"],
                ["./assets/fridge.jpg", "Detailed Caption"],
                ["./assets/marriage.jpg", "Detailed Caption"],
                ["./assets/giraffe.jpg", "Detailed Caption"]
            ],
            inputs=[image_input, style_select]
        )

        # -----------------------------
        # EVENT
        # -----------------------------
        generate_btn.click(
            fn=run_caption,
            inputs=[image_input, style_select],
            outputs=output_text
        )

    return demo