"""
init_knowledge_base.py — скрипт инициализации базы знаний.

Загружает все текстовые документы из указанной директории
в векторное хранилище Qdrant.

Использование:
    python init_knowledge_base.py [директория]

По умолчанию — папка ./docs
"""

import sys
import os
from pathlib import Path

from vector_store import init_collection, add_documents

# Поддерживаемые расширения
TEXT_EXTENSIONS = {".txt", ".md", ".rst", ".csv", ".json"}


def load_directory(docs_dir: str) -> None:
    """
    Рекурсивно обходит директорию и загружает все текстовые файлы в Qdrant.
    """
    base = Path(docs_dir)
    if not base.exists():
        print(f"[init] Директория «{docs_dir}» не найдена. Создаю пример.")
        base.mkdir(parents=True, exist_ok=True)
        # Создаём демонстрационный файл
        sample = base / "sample.txt"
        sample.write_text(
            "LangChain — это фреймворк для построения приложений на основе больших языковых моделей. "
            "Он предоставляет инструменты для цепочек, агентов, памяти и интеграции с внешними источниками данных.\n\n"
            "Qdrant — это высокопроизводительная векторная база данных с открытым исходным кодом, "
            "разработанная специально для хранения и поиска плотных векторов. "
            "Поддерживает фильтрацию по метаданным и масштабируется горизонтально.\n\n"
            "Ollama позволяет запускать большие языковые модели локально на вашем компьютере "
            "без подключения к внешним API. Поддерживает модели llama3, mistral, gemma и другие.",
            encoding="utf-8",
        )
        print(f"[init] Создан демонстрационный файл: {sample}")

    files = [f for f in base.rglob("*") if f.suffix.lower() in TEXT_EXTENSIONS]

    if not files:
        print(f"[init] В директории «{docs_dir}» нет поддерживаемых файлов.")
        return

    print(f"[init] Найдено файлов: {len(files)}")
    print("[init] Инициализация коллекции Qdrant...")
    init_collection()

    total_chunks = 0
    for file in files:
        try:
            content = file.read_text(encoding="utf-8", errors="ignore")
            if not content.strip():
                print(f"  ⚠  Пропущен пустой файл: {file.name}")
                continue
            n = add_documents(content, title=file.name)
            total_chunks += n
            print(f"  ✓  {file.name} → {n} чанков")
        except Exception as e:
            print(f"  ✗  Ошибка при обработке {file.name}: {e}")

    print(f"\n[init] Загрузка завершена. Всего добавлено чанков: {total_chunks}")


if __name__ == "__main__":
    directory = sys.argv[1] if len(sys.argv) > 1 else "./docs"
    load_directory(directory)