import os

import gradio as gr
from PIL import ImageDraw, ImageFont
from PIL.Image import Image


def add_step_counter(image: Image, step_number: int) -> Image:
    """Add a step counter to the bottom left of an image.

    Args:
        image: The PIL Image to modify
        step_number: The step number to display

    Returns:
        A new Image with the step counter added
    """
    img_with_text = image.copy()
    draw = ImageDraw.Draw(img_with_text)

    font = ImageFont.load_default().font_variant(size=24)

    step_text = f"Step {step_number}"
    text_bbox = draw.textbbox((0, 0), step_text, font=font)
    text_width = text_bbox[2] - text_bbox[0]
    text_height = text_bbox[3] - text_bbox[1]

    padding = 8
    x = padding
    y = img_with_text.height - text_height - padding - 2

    draw.rectangle(
        [
            0,
            img_with_text.height - text_height - padding * 2,
            x + text_width + padding * 2,
            img_with_text.height,
        ],
        fill="#d8b4fe",  # Purple color
        outline="#d8b4fe",
    )

    draw.text((x, y), step_text, fill="black", font=font)

    return img_with_text


theme = gr.themes.Base(
    primary_hue="purple",
)


def create_app(self, inputs, prompt, streamer, run_immediately, artifacts_dir):
    with gr.Blocks(theme=theme) as app:
        for input in inputs:
            input.render()

        with gr.Group():
            with gr.Row():
                prompt_box = gr.Textbox(
                    label="Prompt", value=prompt, show_label=False, scale=4
                )
                run_button = gr.Button("Run", variant="primary")
                stop_button = gr.Button("Stop", variant="stop", visible=False)

        chat_log = gr.Chatbot(
            label="Log",
            type="messages",
            group_consecutive_messages=False,
            visible=False,
        )

        @gr.on(
            triggers=[app.load] + [input.change for input in inputs],
            inputs=inputs,
            outputs=[prompt_box],
            trigger_mode="always_last",
        )
        def construct_prompt(*input_values):
            return prompt.format(*input_values)

        run_triggers = [run_button.click]
        if run_immediately:
            run_triggers.append(app.load)

        @gr.on(
            triggers=run_triggers,
            inputs=[prompt_box],
            outputs=[run_button, stop_button, chat_log],
        )
        def run_flow(prompt):
            yield {
                run_button: gr.Button(visible=False),
                chat_log: gr.Chatbot(visible=True),
                stop_button: gr.Button(visible=True),
            }

            log = [gr.ChatMessage(content=prompt, role="user")]
            yield {chat_log: log}

            images_for_gif = []

            for step in streamer(prompt):
                if isinstance(step, str):
                    log.append(gr.ChatMessage(content=step, role="assistant"))
                elif isinstance(step, Image):
                    log.append(gr.ChatMessage(content=gr.Image(step), role="assistant"))

                    img_with_text = add_step_counter(step, len(images_for_gif) + 1)
                    images_for_gif.append(img_with_text)

                    if len(images_for_gif) > 0:
                        gif_path = os.path.join(artifacts_dir, "recording.gif")
                        images_for_gif[0].save(
                            gif_path,
                            save_all=True,
                            append_images=images_for_gif[1:],
                            duration=500,
                            loop=0,
                        )
                elif isinstance(step, gr.ChatMessage):
                    log.append(step)
                else:
                    raise ValueError(
                        f"Cannot handle type: {type(step)} generated by streamer function."
                    )
                yield {chat_log: log}

    return app
