import os
from dotenv import load_dotenv
from langchain_ollama import OllamaLLM
from langchain.agents import initialize_agent, AgentType
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain.tools import tool
from vectorstore import get_vectorstore

# Load environment variables
load_dotenv()
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
if not TAVILY_API_KEY:
    raise ValueError("TAVILY_API_KEY not found in .env file")

# Configuration
PERSIST_DIRECTORY = "./chroma_db"
DOCUMENTS_DIR = "./documents"
EMBEDDING_MODEL = "nomic-embed-text"
LLM_MODEL = "llama3"

# Initialize or load vectorstore
vectorstore = get_vectorstore()

# Define tools
@tool
def search_local_kb(query: str) -> str:
    """Search local knowledge base using ChromaDB."""
    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
    docs = retriever.get_relevant_documents(query)
    return "\n\n".join([doc.page_content for doc in docs])

@tool
def web_search(query: str) -> str:
    """Search the web using Tavily."""
    search = TavilySearchResults(tavily_api_key=TAVILY_API_KEY, max_results=3)
    results = search.run(query)
    return "\n\n".join([result["content"] for result in results])

tools = [search_local_kb, web_search]

# Initialize LLM
llm = OllamaLLM(model=LLM_MODEL)

# Create agent with routing logic
agent = initialize_agent(
    tools=tools,
    llm=llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True,
    handle_parsing_errors=True,
    system_message="""You are an AI agent that must choose between two tools:
    1. search_local_kb: For questions about local documents (concepts, notes, stored information)
    2. web_search: For questions requiring up-to-date information from the internet
    
    When answering, ALWAYS specify your information source:
    - If using local knowledge base: [Source: chromadb]
    - If using web search: [Source: tavily]
    
    Be concise and accurate in your responses."""
)

# Chat interface
def main():
    print("RAG Agent with ChromaDB and Tavily Web Search")
    print("Type 'exit' to quit\n")
    
    while True:
        query = input("Запрос: ").strip()
        if query.lower() == "exit":
            break
        
        if not query:
            continue
            
        try:
            response = agent.run(query)
            print(f"\n{response}\n")
        except Exception as e:
            print(f"Error: {str(e)}\n")

if __name__ == "__main__":
    main()