from pydantic import BaseModel, Field
from typing import Callable, Awaitable, Any, Optional
import re

from apps.webui.models.users import Users
from apps.webui.models.tools import Tools

from utils.misc import get_last_user_message


class Filter:
    class Valves(BaseModel):
        status: bool = Field(default=False)
        match_threshold: int = Field(
            default=1
        )  # Default threshold for number of matching words
        symbols: list = Field(default=["?"])  # Default list of symbols

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

    async def inlet(
        self,
        body: dict,
        __event_emitter__: Callable[[Any], Awaitable[None]],
        __user__: Optional[dict] = None,
        __model__: Optional[dict] = None,
    ) -> dict:
        messages = body["messages"]
        user_message = get_last_user_message(messages)

        if self.valves.status:
            await __event_emitter__(
                {
                    "type": "status",
                    "data": {
                        "description": "Finding the right tools...",
                        "done": False,
                    },
                }
            )

        all_tools = [
            {"id": tool.id, "description": tool.meta.description}
            for tool in Tools.get_tools()
        ]
        available_tool_ids = (
            __model__.get("info", {}).get("meta", {}).get("toolIds", [])
        )
        available_tools = [
            tool for tool in all_tools if tool["id"] in available_tool_ids
        ]

        # Check each word and symbol in user_message against tool descriptions
        user_words = re.findall(r"\w+", user_message.lower())
        user_symbols = [char for char in user_message if char in self.valves.symbols]
        matched_tool_ids = []

        for tool in available_tools:
            description_words = re.findall(r"\w+", tool["description"].lower())
            description_symbols = [
                char for char in tool["description"] if char in self.valves.symbols
            ]
            match_count = sum(
                1 for word in user_words if word in description_words
            ) + sum(1 for symbol in user_symbols if symbol in description_symbols)
            if match_count >= self.valves.match_threshold:
                matched_tool_ids.append(tool["id"])

        if matched_tool_ids:
            body["tool_ids"] = matched_tool_ids
            if self.valves.status:
                await __event_emitter__(
                    {
                        "type": "status",
                        "data": {
                            "description": f"Found matching tools: {', '.join(matched_tool_ids)}",
                            "done": True,
                        },
                    }
                )
            if __event_emitter__:
                await __event_emitter__(
                    {
                        "type": "status",
                        "data": {
                            "description": f"List of tools chosen: {', '.join(matched_tool_ids)}",
                            "done": True,
                        },
                    }
                )
        else:
            if self.valves.status:
                await __event_emitter__(
                    {
                        "type": "status",
                        "data": {
                            "description": "No matching tools found.",
                            "done": True,
                        },
                    }
                )

        return body