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

"""
init_loader.py

Скрипт для загрузки файлов из указанной директории в локальное Qdrant‑хранилище.
Использует:
    * langchain.document_loaders.FileSystemLoader – чтение всех файлов
    * RecursiveCharacterTextSplitter – разбиение на чанки
    * OllamaEmbeddings (nomic-embed-text) – генерация эмбеддингов
    * QdrantVectorStore – сохранение векторных представлений

Параметры:
    --data-dir   Путь к директории с исходными документами
    --collection Название коллекции в Qdrant (по умолчанию: rag_collection)
"""

import argparse
import os
from pathlib import Path
from typing import List, Dict

# LangChain imports
from langchain.document_loaders import DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings.ollama import OllamaEmbeddings
from langchain.vectorstores.qdrant import QdrantVectorStore
from langchain.schema import Document


def load_documents(data_dir: Path) -> List[Document]:
    """
    Загружает все файлы из указанной директории в список объектов Document.
    Поддерживаются форматы *.txt, *.md, *.pdf (если установлен pdfminer).
    """
    loader = DirectoryLoader(
        str(data_dir),
        glob="**/*",
        suffixes=[".txt", ".md", ".pdf"],
        show_progress=True,
    )
    return loader.load()


def chunk_documents(docs: List[Document], chunk_size: int = 1000, overlap: int = 200) -> List[Document]:
    """
    Разбивает каждый документ на чанки фиксированного размера.
    """
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=overlap,
        separators=["\n\n", "\n", " ", ""],
    )
    return splitter.split_documents(docs)


def embed_and_store(chunks: List[Document], collection_name: str) -> None:
    """
    Генерирует эмбеддинги для чанков и сохраняет их в Qdrant.
    """
    # Инициализируем Ollama embeddings
    embedding = OllamaEmbeddings(model="nomic-embed-text")

    # Создаём (или подключаемся к) коллекцию Qdrant
    vector_store = QdrantVectorStore(
        collection_name=collection_name,
        embedding=embedding,
        url="http://localhost:6333",
    )

    # Добавляем документы в хранилище
    vector_store.add_documents(chunks)


def main() -> None:
    parser = argparse.ArgumentParser(description="Загрузка документов в Qdrant.")
    parser.add_argument(
        "--data-dir",
        type=str,
        required=True,
        help="Путь к директории с документами.",
    )
    parser.add_argument(
        "--collection",
        type=str,
        default="rag_collection",
        help="Имя коллекции в Qdrant.",
    )
    args = parser.parse_args()

    data_dir = Path(args.data_dir).expanduser().resolve()
    if not data_dir.is_dir():
        raise FileNotFoundError(f"Каталог {data_dir} не найден.")

    print(f"[INFO] Загружаем документы из: {data_dir}")
    docs = load_documents(data_dir)
    print(f"[INFO] Найдено {len(docs)} документов. Разбиваем на чанки...")

    chunks = chunk_documents(docs)
    print(f"[INFO] Получено {len(chunks)} чанков. Генерируем эмбеддинги и сохраняем в Qdrant...")

    embed_and_store(chunks, args.collection)

    print("[SUCCESS] Загрузка завершена.")


if __name__ == "__main__":
    main()