import os
from pathlib import Path
from typing import Iterable

from langchain_ollama import OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma


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

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

    Returns:
        Chroma: объект vectorstore, готовый к добавлению документов и поиску.
    """
    # Убедимся, что директория для хранения существует
    Path(persist_directory).mkdir(parents=True, exist_ok=True)

    # Создаём эмбеддер Ollama (модель nomic-embed-text)
    embeddings = OllamaEmbeddings(model="nomic-embed-text")

    # Инициализируем Chroma с указанным каталогом хранения
    vectorstore = Chroma(
        persist_directory=persist_directory,
        embedding_function=embeddings,
    )
    return vectorstore


def _load_text_files(directory: str) -> Iterable[str]:
    """
    Генератор, возвращающий содержимое всех .txt и .md файлов из указанной папки.
    Папка может быть вложенной – рекурсивно ищем файлы.

    Args:
        directory (str): Корневая директория для поиска файлов.

    Yields:
        str: Текстовый контент файла.
    """
    for root, _, files in os.walk(directory):
        for file_name in files:
            if file_name.lower().endswith((".txt", ".md")):
                path = Path(root) / file_name
                try:
                    with open(path, "r", encoding="utf-8") as f:
                        yield f.read()
                except Exception as exc:  # pragma: no cover
                    print(f"Не удалось прочитать {path}: {exc}")


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

    Args:
        directory (str): Путь к каталогу, содержащему .txt/.md файлы.
        vectorstore (Chroma): Экземпляр vectorstore, куда будут добавлены документы.
    """
    # Читаем все текстовые файлы из директории
    texts = list(_load_text_files(directory))
    if not texts:
        print(f"В каталоге {directory} не найдено .txt или .md файлов.")
        return

    # Разбиваем каждый документ на чанки
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,  # размер чанка в символах (можно настроить)
        chunk_overlap=200,  # перекрытие для сохранения контекста
    )
    all_chunks = []
    for text in texts:
        chunks = splitter.split_text(text)
        all_chunks.extend(chunks)

    if not all_chunks:
        print("После чанкинга не осталось текста.")
        return

    # Добавляем чанки в vectorstore
    vectorstore.add_texts(all_chunks)

    # Сохраняем изменения (persist)
    vectorstore.persist()
    print(f"Загружено {len(all_chunks)} чанков из {directory} в ChromaDB.")