#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
init_loader.py

Скрипт для загрузки файлов из указанной директории в локальное векторное хранилище Qdrant.
Использует:
    - LangChain OllamaEmbeddings (model "nomic-embed-text")
    - RecursiveCharacterTextSplitter для разбиения на чанки
    - langchain_qdrant.QdrantVectorStore для сохранения эмбеддингов

Параметры запуска:
    python init_loader.py <директория_с_файлами> [--collection <имя_коллекции>]

Пример:
    python init_loader.py ./docs --collection knowledge_base
"""

import argparse
import os
from pathlib import Path
from typing import Iterable

# LangChain imports
from langchain.embeddings.ollama import OllamaEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores.qdrant import QdrantVectorStore
from langchain.document_loaders import TextLoader, UnstructuredFileLoader

# --------------------------------------------------------------------------- #
# Конфигурация по умолчанию
DEFAULT_COLLECTION_NAME = "rag_collection"
QDRANT_HOST = "localhost"
QDRANT_PORT = 6333
EMBEDDING_MODEL = "nomic-embed-text"

# --------------------------------------------------------------------------- #
def _get_text_loaders(directory: Path) -> Iterable[TextLoader]:
    """
    Возвращает генератор загрузчиков для всех файлов в директории.
    Поддерживаются текстовые файлы и любые форматы, которые умеет UnstructuredFileLoader.
    """
    for file_path in directory.rglob("*"):
        if file_path.is_file():
            # Для простоты используем TextLoader для .txt,
            # а остальные обрабатываем через UnstructuredFileLoader
            if file_path.suffix.lower() == ".txt":
                yield TextLoader(str(file_path), encoding="utf-8")
            else:
                yield UnstructuredFileLoader(str(file_path))

# --------------------------------------------------------------------------- #
def load_directory_to_qdrant(
    directory: Path,
    collection_name: str = DEFAULT_COLLECTION_NAME,
) -> None:
    """
    Загружает все документы из указанной директории в Qdrant.

    Args:
        directory (Path): Путь к каталогу с документами.
        collection_name (str): Имя коллекции в Qdrant.
    """
    if not directory.is_dir():
        raise ValueError(f"Путь {directory} не является директорией")

    # 1. Инициализируем эмбеддер
    embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)

    # 2. Создаём текстовый splitter
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
    )

    # 3. Инициализируем QdrantVectorStore (создаст коллекцию при необходимости)
    vector_store = QdrantVectorStore.from_existing_collection(
        embedding=embeddings,
        collection_name=collection_name,
        url=f"http://{QDRANT_HOST}:{QDRANT_PORT}",
    )

    # 4. Загружаем и разбиваем документы, затем сохраняем
    for loader in _get_text_loaders(directory):
        documents = loader.load()
        if not documents:
            continue

        # Разбиваем каждый документ на чанки
        chunks = text_splitter.split_documents(documents)

        # Добавляем векторные представления в Qdrant
        vector_store.add_documents(chunks)

    print(f"Загрузка завершена. Коллекция '{collection_name}' содержит "
          f"{vector_store._client.count(collection_name=collection_name).count} точек.")

# --------------------------------------------------------------------------- #
def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Загрузить документы из директории в Qdrant."
    )
    parser.add_argument(
        "directory",
        type=str,
        help="Путь к каталогу с документами.",
    )
    parser.add_argument(
        "--collection",
        type=str,
        default=DEFAULT_COLLECTION_NAME,
        help=f"Имя коллекции (по умолчанию: {DEFAULT_COLLECTION_NAME}).",
    )
    return parser.parse_args()

# --------------------------------------------------------------------------- #
if __name__ == "__main__":
    args = _parse_args()
    load_directory_to_qdrant(Path(args.directory), collection_name=args.collection)