import os
from dotenv import load_dotenv
from langchain_ollama import OllamaLLM, OllamaEmbeddings
from langchain_chroma import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import DirectoryLoader, TextLoader, UnstructuredMarkdownLoader
from langchain.tools import Tool
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain.agents import initialize_agent, AgentType

# 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"

def setup_vectorstore():
    """Initialize or load ChromaDB vectorstore with Ollama embeddings."""
    embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
    
    if os.path.exists(PERSIST_DIRECTORY) and os.listdir(PERSIST_DIRECTORY):
        vectorstore = Chroma(
            persist_directory=PERSIST_DIRECTORY,
            embedding_function=embeddings
        )
    else:
        # Load and process documents
        loader = DirectoryLoader(
            DOCUMENTS_DIR,
            glob="**/*",
            loader_cls=lambda path: TextLoader(path, encoding="utf-8") if path.endswith(".txt") 
                                    else UnstructuredMarkdownLoader(path) if path.endswith(".md") 
                                    else None,
            show_progress=True,
            use_multithreading=True
        )
        documents = loader.load()
        
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            length_function=len
        )
        texts = text_splitter.split_documents(documents)
        
        vectorstore = Chroma.from_documents(
            documents=texts,
            embedding=embeddings,
            persist_directory=PERSIST_DIRECTORY
        )
        vectorstore.persist()
    
    return vectorstore

# Initialize vectorstore
vectorstore = setup_vectorstore()

# Define tools
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])

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])

# Create LangChain tools
local_tool = Tool(
    name="search_local_kb",
    func=search_local_kb,
    description="Useful for answering questions about local documents stored in the knowledge base."
)

web_tool = Tool(
    name="web_search",
    func=web_search,
    description="Useful for answering questions about current events, news, or general knowledge from the internet."
)

tools = [local_tool, web_tool]

# 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()