"""Tools used by the RAG agent.

Two tools are provided:

* :func:`search_local_kb` – semantic search in the local ChromaDB vector store.
* :func:`web_search` – real‑time web search via Tavily.

Both functions are decorated with :func:`langchain.tools.tool` so that LangChain
can expose them to the agent.
"""

from __future__ import annotations

import os
from typing import List

from langchain.chains import RetrievalQA
from langchain.chroma import Chroma
from langchain_ollama import ChatOllama
from langchain_tavily import TavilySearchResults
from langchain.tools import tool

# Load environment variables (TAVILY_API_KEY)
from dotenv import load_dotenv

load_dotenv()

# ---------------------------------------------------------------------------
# Local KB search
# ---------------------------------------------------------------------------
@tool("search_local_kb")
def search_local_kb(query: str, top_k: int = 3, vectorstore: Chroma | None = None) -> str:
    """Return the best answer from the local ChromaDB store.

    Parameters
    ----------
    query: str
        The user question.
    top_k: int
        Number of documents to retrieve.
    vectorstore: Chroma | None
        If ``None`` the function will attempt to load the default store.

    Returns
    -------
    str
        The answer prefixed with the source identifier.
    """
    if vectorstore is None:
        # Load the default store
        from vectorstore import create_vectorstore
        vectorstore = create_vectorstore()

    retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
    qa_chain = RetrievalQA.from_chain_type(
        llm=ChatOllama(model="llama3", temperature=0),
        chain_type="stuff",
        retriever=retriever,
    )
    result = qa_chain.run(query)
    return f"[chromadb] {result}"

# ---------------------------------------------------------------------------
# Web search via Tavily
# ---------------------------------------------------------------------------
@tool("web_search")
def web_search(query: str, top_k: int = 3) -> str:
    """Return an answer based on a Tavily web search.

    Parameters
    ----------
    query: str
        The user question.
    top_k: int
        Number of search results to retrieve.

    Returns
    -------
    str
        The answer prefixed with the source identifier.
    """
    api_key = os.getenv("TAVILY_API_KEY")
    if not api_key:
        raise RuntimeError("TAVILY_API_KEY environment variable is not set")
    tavily = TavilySearchResults(api_key=api_key, top_k=top_k)
    results = tavily.run(query)
    snippets = "\n\n".join(item.get("content", "") for item in results)
    llm = ChatOllama(model="llama3", temperature=0)
    prompt = f"Answer the question based on the following web snippets:\n{snippets}\n\nQuestion: {query}\nAnswer:"  # noqa: E501
    answer = llm.invoke(prompt).content
    return f"[tavily] {answer}"