"""Tools for the RAG agent: local KB search and web search via Tavily."""

from typing import List, Dict

from langchain_ollama import ChatOllama
from langchain_chroma import Chroma
from langchain_tavily import TavilySearchResults
from langchain.tools import tool

# --- Local KB search tool -----------------------------------------------------

@tool("search_local_kb")
def search_local_kb(query: str, top_k: int = 3, vectorstore: Chroma = None) -> List[Dict]:
    """Perform a semantic search in the local Chroma vector store.

    Parameters
    ----------
    query: str
        The user query.
    top_k: int
        Number of top results to return.
    vectorstore: Chroma
        The vector store to search.

    Returns
    -------
    List[Dict]
        List of dictionaries containing ``content`` and ``metadata``.
    """
    if vectorstore is None:
        raise ValueError("vectorstore must be provided")
    retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
    docs = retriever.get_relevant_documents(query)
    return [{"content": doc.page_content, "metadata": doc.metadata} for doc in docs]

# --- Web search tool ---------------------------------------------------------

@tool("web_search")
def web_search(query: str, top_k: int = 3) -> List[Dict]:
    """Search the web using Tavily.

    Parameters
    ----------
    query: str
        The user query.
    top_k: int
        Number of top results to return.

    Returns
    -------
    List[Dict]
        List of dictionaries containing ``title``, ``url`` and ``content``.
    """
    tavily = TavilySearchResults(max_results=top_k)
    results = tavily.run(query)
    # Tavily returns a list of dicts with keys: title, url, content
    return results