#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Скрипт для создания коллекции в Qdrant на основе core.conf.
Используется для единой коллекции всех ботов.
"""

import sys
import argparse
import yaml
from pathlib import Path

from qdrant_client import QdrantClient
from qdrant_client.http import models


def load_config(config_path: str) -> dict:
    """Загружает core.conf как YAML."""
    with open(config_path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)


def main():
    parser = argparse.ArgumentParser(
        description="Создать коллекцию в Qdrant для RAG-сервера"
    )
    parser.add_argument(
        "--config",
        default="core/core.conf",
        help="Путь к core.conf (по умолчанию core/core.conf)"
    )
    parser.add_argument(
        "--collection",
        help="Имя коллекции (переопределяет значение из core.conf)"
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Удалить существующую коллекцию, если она есть, и создать заново"
    )
    args = parser.parse_args()

    # Загружаем конфигурацию
    config_path = Path(args.config)
    if not config_path.exists():
        print(f"Ошибка: файл {config_path} не найден")
        sys.exit(1)

    try:
        config = load_config(config_path)
    except Exception as e:
        print(f"Ошибка загрузки конфигурации: {e}")
        sys.exit(1)

    # Извлекаем параметры Qdrant
    qdrant_cfg = config.get('qdrant')
    if not qdrant_cfg:
        print("Ошибка: в конфиге отсутствует секция 'qdrant'")
        sys.exit(1)

    host = qdrant_cfg.get('host', 'localhost')
    port = qdrant_cfg.get('port', 6333)
    vector_size = qdrant_cfg.get('vector_size', 1024)
    distance = qdrant_cfg.get('distance', 'Cosine')
    collection = args.collection or qdrant_cfg.get('collection')

    if not collection:
        print("Ошибка: не указано имя коллекции (в core.conf или через --collection)")
        sys.exit(1)

    print(f"Подключение к Qdrant: {host}:{port}")
    client = QdrantClient(host=host, port=port)

    # Проверяем существование коллекции
    if client.collection_exists(collection):
        if args.force:
            print(f"Коллекция '{collection}' существует. Удаляем...")
            client.delete_collection(collection)
        else:
            print(f"Коллекция '{collection}' уже существует. Используйте --force для пересоздания.")
            sys.exit(0)

    # Создаём коллекцию
    try:
        client.create_collection(
            collection_name=collection,
            vectors_config=models.VectorParams(
                size=vector_size,
                distance=models.Distance[distance.upper()]
            )
        )
        print(f"Коллекция '{collection}' успешно создана (size={vector_size}, distance={distance})")
    except Exception as e:
        print(f"Ошибка создания коллекции: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()