from typing import List

from langchain.tools import tool
from langchain_ollama import Ollama
from langchain_qdrant import QdrantVectorStore
from langchain_tavily import TavilySearchResults


@tool
def search_local_kb(
    query: str,
    top_k: int,
    vectorstore: QdrantVectorStore,
) -> List[str]:
    """
    Perform a semantic search in the local knowledge base stored in Qdrant.

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

    Returns
    -------
    List[str]
        List of relevant document snippets.
    """
    retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
    docs = retriever.get_relevant_documents(query)
    return [doc.page_content for doc in docs]


@tool
def web_search(
    query: str,
    tavily_api_key: str,
    max_results: int = 3,
) -> List[str]:
    """
    Perform a web search using Tavily.

    Parameters
    ----------
    query : str
        The user's query.
    tavily_api_key : str
        Tavily API key.
    max_results : int
        Number of search results to return.

    Returns
    -------
    List[str]
        List of search result snippets.
    """
    tavily = TavilySearchResults(
        api_key=tavily_api_key,
        max_results=max_results,
    )
    results = tavily.run(query)
    # Extract snippets from results
    snippets = []
    for result in results:
        snippet = result.get("content") or result.get("snippet") or result.get("title")
        if snippet:
            snippets.append(snippet)
    return snippets