from typing import Any, Dict, List

from openrouter.components import ChatToolMessageContentTypedDict
from tavily import TavilyClient

from ai.tool import Tool
from ai.utils import *


class TavilySearchTool(Tool):
    def __init__(self, tavily_token: str):
        self._client = TavilyClient(api_key=tavily_token)
    
    @property
    def name(self) -> str:
        return "tavily_search"
    
    @property
    def description(self) -> str:
        return "Поиск информации в интернете"
    
    @property
    def parameters(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Запрос для поиска (на русском или английском языке)"
                }
            },
            "required": ["query"]
        }
    
    async def execute(self, args: Dict[str, Any], _artifacts: Dict[str, Any]) -> List[ChatToolMessageContentTypedDict]:
        query = args.get("query", "")
        print(f"Веб-поиск: {query}")
        
        try:
            results = self._client.search(query=query, max_results=5)
            
            if not results or "results" not in results:
                return serialize_message_content(text="Не удалось получить результаты поиска.")
            
            answer_parts = []
            for i, result in enumerate(results["results"], 1):
                title = result.get("title", "Без названия")
                url = result.get("url", "")
                content = result.get("content", "")
                answer_parts.append(f"{i}. {title}\n   {url}\n   {content}\n")
            
            answer = "\n".join(answer_parts)
            return serialize_message_content(text=f"По запросу \"{query}\" найдено:\n\n{answer}")
        except Exception as e:
            print(f"Ошибка веб-поиска: {e}")
            return serialize_message_content(text=f"Не удалось выполнить веб-поиск: {e}")