"""ChromaDB + OllamaEmbeddings (nomic-embed-text)."""
from pathlib import Path
from langchain_ollama import OllamaEmbeddings
from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document


def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
    """Создать/открыть ChromaDB с Ollama-эмбеддингами."""
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    return Chroma(
        collection_name="rag_kb",
        embedding_function=embeddings,
        persist_directory=persist_directory,
    )


def load_documents(directory: str, vectorstore: Chroma) -> int:
    """Загрузить .txt/.md файлы из директории в ChromaDB с чанкингом."""
    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    docs = []
    for path in Path(directory).glob("**/*"):
        if path.suffix.lower() not in (".txt", ".md"):
            continue
        text = path.read_text(encoding="utf-8", errors="replace")
        chunks = splitter.split_text(text)
        docs.extend(
            [Document(page_content=c, metadata={"source": str(path)}) for c in chunks]
        )
    if docs:
        vectorstore.add_documents(docs)
    return len(docs)