"""
title: Dynamic Vision Router
author: open-webui, atgehrhardt
author_url: https://github.com/open-webui
funding_url: https://github.com/open-webui
version: 0.1.4
required_open_webui_version: 0.3.8
"""

from pydantic import BaseModel, Field
from typing import Callable, Awaitable, Any, Optional, Literal
import json

from open_webui.utils.misc import get_last_user_message_item


class Filter:
    class Valves(BaseModel):
        vision_model_id: str = Field(
            default="",
            description="The identifier of the vision model to be used for processing images. Note: Compatibility is provider-specific; ollama models can only route to ollama models, and OpenAI models to OpenAI models respectively.",
        )
        status: bool = Field(
            default=False,
            description="A flag to enable or disable the status indicator. Set to True to enable status updates.",
        )
        pass

    def __init__(self):
        self.valves = self.Valves()
        pass

    async def inlet(
        self,
        body: dict,
        __event_emitter__: Callable[[Any], Awaitable[None]],
        __model__: Optional[dict] = None,
    ) -> dict:
        if __model__["id"] == self.valves.vision_model_id:
            return body

        messages = body["messages"]
        user_message = get_last_user_message_item(messages)

        has_images = "images" in user_message or (
            isinstance(user_message.get("content"), list)
            and any(item.get("type") == "image_url" for item in user_message["content"])
        )

        if has_images:
            if self.valves.vision_model_id:
                body["model"] = self.valves.vision_model_id
                if self.valves.status:
                    await __event_emitter__(
                        {
                            "type": "status",
                            "data": {
                                "description": f"Request routed to {self.valves.vision_model_id}",
                                "done": True,
                            },
                        }
                    )
            else:
                if self.valves.status:
                    await __event_emitter__(
                        {
                            "type": "status",
                            "data": {
                                "description": "No vision model ID provided, routing could not be completed.",
                                "done": True,
                            },
                        }
                    )

        return body