# -------------------- vectorstore.py --------------------
from pathlib import Path
from langchain_chroma import Chroma
from langchain_ollama import OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

def create_vectorstore(persist_directory: str = "./chroma_db"):
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    vector_store = Chroma(
        collection_name="rag_collection",
        embedding_function=embeddings,
        persist_directory=persist_directory,
    )
    return vector_store

def load_documents(directory: str, vectorstore):
    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    docs = []
    for file_path in Path(directory).glob("*.txt"):
        text = file_path.read_text(encoding="utf-8")
        docs.extend(splitter.create_documents([text]))
    for file_path in Path(directory).glob("*.md"):
        text = file_path.read_text(encoding="utf-8")
        docs.extend(splitter.create_documents([text]))
    vectorstore.add_documents(docs)

# -------------------- tools.py --------------------
from langchain.tools import tool
from langchain_ollama import ChatOllama

@tool
def search_local_kb(query: str, top_k: int = 3) -> str:
    """Semantic search in the local ChromaDB knowledge base."""
    retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
    docs = retriever.invoke({"query": query})["documents"]
    return "\n".join(doc.page_content for doc in docs)

@tool
def web_search(query: str) -> str:
    """Web search using Tavily."""
    from langchain_tavily import TavilySearchResults
    tavily = TavilySearchResults(api_key=__import__("os").environ["TAVILY_API_KEY"])
    results = tavily.invoke({"query": query})
    return "\n".join(f"{r['title']}: {r['url']}" for r in results)

# -------------------- agent.py --------------------
from langchain.agents import create_agent
from langchain_ollama import ChatOllama

llm = ChatOllama(model="llama3", temperature=0.2)

system_prompt = """
You are an assistant that answers user questions.
If the answer can be found in the local knowledge base, use `search_local_kb`.
Otherwise, use `web_search`.  
Always indicate the source of your answer: either "chromadb" or "tavily".
"""

agent = create_agent(
    model=llm,
    tools=[search_local_kb, web_search],
    system_prompt=system_prompt,
)

# -------------------- main.py --------------------
import os
from dotenv import load_dotenv

load_dotenv()

if __name__ == "__main__":
    # Initialize vectorstore and load documents if not already loaded
    vectorstore = create_vectorstore()
    if not vectorstore.get_collection().count():
        load_documents("documents", vectorstore)
        vectorstore.persist()

    print("Chat started. Type 'exit' to quit.")
    while True:
        user_input = input("\nЗапрос: ").strip()
        if user_input.lower() in ("exit", "quit", "выход"):
            break
        result = agent.invoke({"messages": [{"role": "human", "content": user_input}]})
        ai_msg = result["messages"][-1]
        print(f"[{ai_msg.tool_calls[0]['name'].capitalize()}] {ai_msg.content}")
        source = "chromadb" if ai_msg.tool_calls[0]["name"] == "search_local_kb" else "tavily"
        print(f"Источник: {source}")