from typing import List

import gradio as gr
import PIL
import spaces
from gradio import ChatMessage
from smolagents.gradio_ui import stream_to_gradio

from agents.all_agents import get_master_agent
from llm import TRANSFORMERS_MODEL_IDS, get_transformers_model


def resize_image(image):
    width, height = image.size
    if width > 1200 or height > 800:
        ratio = min(1200 / width, 800 / height)
        new_width = int(width * ratio)
        new_height = int(height * ratio)
        resized_image = image.resize((new_width, new_height), PIL.Image.Resampling.LANCZOS)
        return resized_image
    return image


@spaces.GPU(duration=300)
def chat_interface_fn(input_request, history: List[ChatMessage], gallery, transformers_model_id):
    model = get_transformers_model(transformers_model_id)
    agent = get_master_agent(model)
    if gallery is None:
        gallery = []
    else:
        gallery = [value[0] for value in gallery]
    message = input_request["text"]
    image_paths = input_request["files"]
    prompt = f"""
    You are given the following message from the user:
    {message}
    """

    if len(image_paths) > 0:
        prompt += """
        The user also provided the additional images that you can find in "images" variable
        """
    if len(history) > 0:
        prompt += "This request follows a previous request, you can use the previous request to help you answer the current request."

    prompt += """ 
    Before your final answer, if you have any images to show, store them in the "final_images" variable.
    Always return a text of what you did.

    Never assume an invented model name, always use the model name provided by the task_model_retriever tool.
    """

    images = [PIL.Image.open(image_path) for image_path in image_paths]
    if len(gallery) > 0:
        images.extend(gallery)
    resized_images = [resize_image(image) for image in images]

    for message in stream_to_gradio(
        agent,
        task=prompt,
        task_images=resized_images,
        additional_args={"images": images},
        reset_agent_memory=False,
    ):
        history.append(message)
        yield history, None

    final_images = agent.python_executor.state.get("final_images", [])
    gallery.extend(final_images)
    yield history, gallery


def example_selected(example):
    textbox.value = example[0]
    image_box.value = example[1]

    example = {
        "text": example[0],
        "files": [
            {
                "url": example[1],
                "path": example[1],
                "name": example[1],
            }
        ],
    }

    return example


with gr.Blocks() as demo:
    gr.Markdown(
        """
    # ScouterAI
    """
    )
    gr.HTML(
        """
        <div style="display: flex; align-items: center; gap: 20px; margin: 20px 0;">
            <img src="https://cdn-uploads.huggingface.co/production/uploads/632885ba1558dac67c440aa8/KpMuW4Qvrh5N-FMcVKKqG.png" 
            alt="Picture" 
            style="max-height: 350px; flex-shrink: 0;" />
            <div style="flex-grow: 1;">
                <p style="margin: 0; font-size: 1.1em;">
                    <p style="font-size: 1.8em; margin-bottom: 10px; font-weight: bold">Welcome to ScouterAI</p>
                    <p style="font-size: 1.2em;">The agent capable of identifying the best 
                    model among the entire HuggingFace Hub to use for your needs.</p>
                    This Space focuses on using agentic reasoning to plan the use of multiple models to perform vision tasks.
                    <br>
                    To answer your request, the agent will use the following models from the hub:
                    <br>
                    <ul>
                        <li><a href="https://huggingface.co/models?pipeline_tag=object-detection&library=transformers&sort=trending">Object detection</a></li>
                        <li><a href="https://huggingface.co/models?pipeline_tag=image-segmentation&library=transformers&sort=trending">Image segmentation</a></li>
                        <li><a href="https://huggingface.co/models?pipeline_tag=image-classification&library=transformers&sort=trending">Image classification</a></li>
                    </ul>
                    The agent can resize and crop images as well as annotating it with bounding boxes, masks and labels.
                    <br>
                    <br>
                    Type your request and add images to the textbox below or click on one of the examples to see how <strong style="font-size: 1.5em;">powerful</strong> it is.
                </p>
            </div>
        </div>
        """,
    )
    gr.Markdown(
        """
        ## Update 17/06/2025
        Welcome to the new version of ScouterAI!<br>
        This Space is now powered by ZeroGPU, which means you can use it for free!<br>
        You can now use any model from the HuggingFace Hub, just provide the model ID below.<br>
        <br>
        """
    )
    transformers_model_id = gr.Dropdown(label="Transformers Model ID", info="Select a model to use or provide your own model ID", choices=TRANSFORMERS_MODEL_IDS)
    output_gallery = gr.Gallery(label="Images generated by the agent (do not put images)", type="pil", format="png")
    textbox = gr.MultimodalTextbox()
    gr.ChatInterface(
        chat_interface_fn,
        type="messages",
        multimodal=True,
        textbox=textbox,
        additional_inputs=[output_gallery, transformers_model_id],
        additional_outputs=[output_gallery],
    )

    text_box = gr.Textbox(label="Text", visible=False)
    image_box = gr.Image(label="Image", visible=False)
    dataset = gr.Dataset(
        samples=[
            [
                "I would like to detect all the cars in the image",
                "https://upload.wikimedia.org/wikipedia/commons/5/51/Crossing_the_Hudson_River_on_the_George_Washington_Bridge_from_Fort_Lee%2C_New_Jersey_to_Manhattan%2C_New_York_%287237796950%29.jpg",
            ],
            [
                "Find vegetables in the image and annotate the image with their masks",
                "https://media.istockphoto.com/id/1203599923/fr/photo/fond-de-nourriture-avec-lassortiment-des-l%C3%A9gumes-organiques-frais.jpg?s=612x612&w=0&k=20&c=Yu8nfOYI9YZ0UTpb7iFqX8OHp9wfvd9keMQ0BZIzhWs=",
            ],
            [
                "Detect each dog in the image and identify its breed, then provide a crop of each dog and annotate the original image with a bounding box and a label",
                "https://images.pexels.com/photos/10094979/pexels-photo-10094979.jpeg",
            ],
        ],
        components=[text_box, image_box],
        label="Examples",
    )

    dataset.select(example_selected, [dataset], [textbox])


demo.launch()