import os
from langchain_community.tools.tavily import TavilySearchResults
from langchain_core.tools import tool
from vectorstore import create_vectorstore


# Global vectorstore instance
_vectorstore = None


def get_vectorstore():
    """Get or create the global vectorstore instance."""
    global _vectorstore
    if _vectorstore is None:
        _vectorstore = create_vectorstore()
    return _vectorstore


@tool
def search_local_kb(query: str, top_k: int = 5) -> str:
    """Search for relevant information in the local knowledge base (ChromaDB).
    
    Use this tool for questions about local documents, notes, or stored knowledge.
    
    Args:
        query: The search query
        top_k: Number of top results to return (default: 5)
    
    Returns:
        Search results from the local knowledge base
    """
    vectorstore = get_vectorstore()
    retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
    docs = retriever.invoke(query)
    
    if not docs:
        return "No relevant information found in local knowledge base."
    
    results = []
    for i, doc in enumerate(docs, 1):
        source = doc.metadata.get("source", "unknown")
        results.append(f"[Document {i}] {doc.page_content}\n(Source: {source})")
    
    return "\n\n".join(results)


@tool
def web_search(query: str) -> str:
    """Search the web for current information using Tavily.
    
    Use this tool for questions about current events, recent news, or facts 
    that may not be in the local knowledge base.
    
    Args:
        query: The search query
    
    Returns:
        Web search results
    """
    tavily_api_key = os.getenv("TAVILY_API_KEY")
    if not tavily_api_key:
        return "Error: TAVILY_API_KEY not set in environment"
    
    search = TavilySearchResults(
        max_results=5,
        api_key=tavily_api_key
    )
    results = search.invoke(query)
    
    if not results:
        return "No web search results found."
    
    formatted_results = []
    for i, result in enumerate(results, 1):
        title = result.get("title", "No title")
        content = result.get("content", "No content")
        url = result.get("url", "No URL")
        formatted_results.append(f"[Result {i}] {title}\n{content}\nURL: {url}")
    
    return "\n\n".join(formatted_results)