import os
from typing import Any

from langchain.tools import tool
from langchain_tavily import TavilySearch

from vectorstore import create_vectorstore


def _format_local_docs(results: list[Any]) -> str:
    if not results:
        return "Источник: chromadb\nНет релевантных локальных документов."

    parts = ["Источник: chromadb"]
    for index, doc in enumerate(results, start=1):
        source = doc.metadata.get("source", "неизвестно")
        parts.append(f"{index}. ({source}) {doc.page_content}")
    return "\n".join(parts)


@tool
def search_local_kb(query: str, top_k: int = 3) -> str:
    """Search the local ChromaDB knowledge base for internal course notes and local documents."""
    collection = create_vectorstore()
    retriever = collection.as_retriever(search_kwargs={"k": top_k})
    results = retriever.invoke(query)
    return _format_local_docs(results)


@tool
def web_search(query: str) -> str:
    """Search the web with Tavily for fresh facts, current news, or anything not covered by local documents."""
    api_key = os.getenv("TAVILY_API_KEY")
    if not api_key:
        return "Источник: tavily\nTAVILY_API_KEY не задан, поиск в интернете недоступен."

    search_tool = TavilySearch(
        max_results=5,
        topic="general",
    )
    result = search_tool.invoke({"query": query})
    return f"Источник: tavily\n{result}"