# 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") -> Chroma:
    """Create a persistent Chroma vector store with Ollama embeddings."""
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    return Chroma(
        collection_name="rag_collection",
        embedding_function=embeddings,
        persist_directory=persist_directory,
    )


def load_documents(directory: str, vectorstore: Chroma) -> None:
    """Load .txt and .md files from `directory`, chunk them, and add to the store."""
    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    docs = []
    for file in Path(directory).glob("*.txt") | Path(directory).glob("*.md"):
        text = file.read_text(encoding="utf-8")
        docs.extend(splitter.create_documents([text]))
    vectorstore.add_documents(docs)


# tools.py
from langchain.tools import tool

from langchain_chroma import Chroma
from langchain_ollama import OllamaEmbeddings


@tool("search_local_kb")
def search_local_kb(query: str, top_k: int = 3) -> str:
    """Semantic search in local knowledge base."""
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    vectorstore = Chroma(
        collection_name="rag_collection",
        embedding_function=embeddings,
        persist_directory="./chroma_db",
    )
    retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
    docs = retriever.invoke(query)
    return "\n".join(doc.page_content for doc in docs)


@tool("web_search")
def web_search(query: str) -> str:
    """Search the web using Tavily."""
    from langchain_tavily import TavilyAPIWrapper

    tavily = TavilyAPIWrapper()
    results = tavily.run(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

from tools import search_local_kb, web_search


def create_rag_agent():
    llm = ChatOllama(model="llama3", temperature=0.2)
    tools = [search_local_kb, web_search]
    system_prompt = (
        "You are an assistant that answers user questions.\n"
        "If the answer can be found in local documents, use search_local_kb.\n"
        "Otherwise, use web_search. Return the answer followed by a line\n"
        "\"Source: chromadb\" or \"Source: tavily\"."
    )
    return create_agent(model=llm, tools=tools, system_prompt=system_prompt)


# main.py
import os
from pathlib import Path

from dotenv import load_dotenv

from vectorstore import create_vectorstore, load_documents
from agent import create_rag_agent


def init_db():
    """Create and populate the Chroma DB if it does not exist."""
    db_path = Path("./chroma_db")
    if not db_path.exists():
        store = create_vectorstore()
        load_documents("documents", store)
        store.persist()


if __name__ == "__main__":
    load_dotenv()  # Loads TAVILY_API_KEY and any other env vars
    init_db()
    agent = create_rag_agent()

    while True:
        user_input = input("Запрос: ").strip()
        if not user_input or user_input.lower() in ("exit", "quit", "выход"):
            break

        result = agent.invoke({"messages": [{"role": "human", "content": user_input}]})
        for msg in result["messages"]:
            # `msg` is a Pydantic model; use the `.content` attribute
            if hasattr(msg, "content"):
                print(msg.content)