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

from typing import List

from langchain_community.tools.tavily import TavilySearchResults
from langchain.tools import tool
from langchain_chroma import Chroma

# ---------------------------------------------------------------------------
# Local KB search tool
# ---------------------------------------------------------------------------
@tool("search_local_kb")

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

    Parameters
    ----------
    query: str
        The user query.
    top_k: int, optional
        Number of top results to return.
    vectorstore: Chroma, optional
        The vector store to query. If None, the function will raise an error.

    Returns
    -------
    str
        Concatenated content of the top results.
    """
    if vectorstore is None:
        raise ValueError("vectorstore must be provided to search_local_kb")
    retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
    docs = retriever.get_relevant_documents(query)
    return "\n\n---\n\n".join(doc.page_content for doc in docs)

# ---------------------------------------------------------------------------
# Web search tool using Tavily
# ---------------------------------------------------------------------------
@tool("web_search")

def web_search(query: str, top_k: int = 3) -> str:
    """Search the web via Tavily and return a formatted string of results.

    Parameters
    ----------
    query: str
        The search query.
    top_k: int, optional
        Number of top results to return.

    Returns
    -------
    str
        Formatted search results.
    """
    tavily = TavilySearchResults(max_results=top_k)
    results = tavily.run(query)
    formatted = []
    for i, r in enumerate(results, 1):
        formatted.append(f"{i}. {r.get('title', 'No title')}\n{r.get('url', '')}\n{r.get('content', '')}")
    return "\n\n---\n\n".join(formatted)