# vectorstore.py
"""
Модуль для работы с ChromaDB и OllamaEmbeddings.
"""

import os
from pathlib import Path
from typing import Iterable

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


def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
    """
    Создаёт (или загружает) коллекцию ChromaDB с OllamaEmbeddings.

    Args:
        persist_directory: Путь, где будет храниться база данных.
                            Если директория не существует – создаётся автоматически.

    Returns:
        Chroma объект, готовый к работе.
    """
    # Создаём путь и при необходимости создаём папку
    path = Path(persist_directory)
    path.mkdir(parents=True, exist_ok=True)

    # Инициализируем эмбеддер Ollama
    embeddings = OllamaEmbeddings(model="nomic-embed-text")

    # Создаём/загружаем коллекцию Chroma
    vectorstore = Chroma(
        persist_directory=str(path),
        embedding_function=embeddings,
    )
    return vectorstore


def load_documents(directory: str | Path, vectorstore: Chroma) -> None:
    """
    Читает все .txt и .md файлы из указанной директории, разбивает их на чанки
    и добавляет в коллекцию ChromaDB.

    Args:
        directory: Путь к папке с документами.
        vectorstore: Объект Chroma, куда будут загружены документы.
    """
    # Подготовка текстового разделителя
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
    )

    docs_to_add: list[Document] = []

    # Ищем файлы .txt и .md рекурсивно
    for file_path in Path(directory).rglob("*"):
        if not file_path.is_file():
            continue
        if file_path.suffix.lower() not in {".txt", ".md"}:
            continue

        try:
            text = file_path.read_text(encoding="utf-8")
        except Exception as exc:
            # Если файл не читается – пропускаем и выводим предупреждение
            print(f"⚠️ Не удалось прочитать {file_path}: {exc}")
            continue

        # Разбиваем текст на чанки
        chunks = splitter.split_text(text)

        # Создаём Document объекты с метаданными
        for i, chunk in enumerate(chunks):
            docs_to_add.append(
                Document(
                    page_content=chunk,
                    metadata={
                        "source": str(file_path),
                        "page": i + 1,
                    },
                )
            )

    if not docs_to_add:
        print("⚠️ В указанной папке не найдено подходящих файлов.")
        return

    # Добавляем документы в коллекцию
    vectorstore.add_documents(docs_to_add)

    # Сохраняем изменения (persist)
    vectorstore.persist()