"""ChromaDB + Ollama embeddings: создание хранилища и загрузка документов."""
from __future__ import annotations

from pathlib import Path

from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_ollama import OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

DEFAULT_PERSIST_DIR = "./chroma_db"
DEFAULT_EMBED_MODEL = "nomic-embed-text"
CHUNK_SIZE = 800
CHUNK_OVERLAP = 120


def create_vectorstore(persist_directory: str = DEFAULT_PERSIST_DIR) -> Chroma:
    embeddings = OllamaEmbeddings(model=DEFAULT_EMBED_MODEL)
    return Chroma(
        collection_name="local_kb",
        embedding_function=embeddings,
        persist_directory=persist_directory,
    )


def _read_text_file(path: Path) -> str:
    return path.read_text(encoding="utf-8")


def load_documents(directory: str, vectorstore: Chroma) -> int:
    """Читает .txt/.md из directory, чанкует и добавляет в ChromaDB."""
    root = Path(directory)
    if not root.exists():
        return 0

    files = sorted(root.glob("*.txt")) + sorted(root.glob("*.md"))
    if not files:
        return 0

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
    )
    docs: list[Document] = []
    for fp in files:
        text = _read_text_file(fp)
        chunks = splitter.split_text(text)
        for i, chunk in enumerate(chunks):
            docs.append(
                Document(
                    page_content=chunk,
                    metadata={"source": fp.name, "chunk": i},
                )
            )

    if not docs:
        return 0

    vectorstore.add_documents(docs)
    return len(docs)