import os
from pathlib import Path
from typing import List

from langchain_ollama import OllamaEmbeddings
from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.docstore.document import Document


def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
    """
    Создаёт и возвращает объект Chroma, подключённый к указанной директории.
    Если директория не существует, она будет создана.
    """
    Path(persist_directory).mkdir(parents=True, exist_ok=True)
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    return Chroma(
        persist_directory=persist_directory,
        embedding_function=embeddings,
    )


def _load_text_files(directory: str) -> List[Document]:
    """
    Читает все .txt и .md файлы из указанной директории и возвращает список Document.
    """
    docs: List[Document] = []
    for file_path in Path(directory).rglob("*"):
        if file_path.suffix.lower() in {".txt", ".md"}:
            text = file_path.read_text(encoding="utf-8")
            docs.append(Document(page_content=text, metadata={"source": str(file_path)}))
    return docs


def load_documents(directory: str, vectorstore: Chroma) -> None:
    """
    Загружает документы из указанной директории в ChromaDB.
    Делает чанкинг с помощью RecursiveCharacterTextSplitter и добавляет в коллекцию.
    """
    # Читаем файлы
    raw_docs = _load_text_files(directory)

    # Разбиваем на чанки
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        separators=["\n\n", "\n", " ", ""],
    )
    chunks = splitter.split_documents(raw_docs)

    # Добавляем в коллекцию
    vectorstore.add_documents(chunks)
    # Сохраняем изменения
    vectorstore.persist()
    print(f"Загружено {len(chunks)} чанков из {len(raw_docs)} документов в {vectorstore.persist_directory}")