# qdrant_client.py
"""
Модуль инициализации клиента Qdrant и простых операций с ним.
Используется в агенте для RAG‑памяти.

Поддержка:
- Создание/получение коллекции
- Добавление документов (с чанками) с эмбеддингами Ollama
- Семантический поиск по запросу

Требуемые зависимости:
    pip install langchain langchain-qdrant langchain-ollama qdrant-client
"""

from __future__ import annotations

import os
from pathlib import Path
from typing import Iterable, List, Tuple

from langchain_ollama import OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter
from qdrant_client import QdrantClient as _QdrantClient
from qdrant_client.http.models import Distance, VectorParams

# --------------------------------------------------------------------------- #
# Конфигурация (можно переопределить через переменные окружения)
# --------------------------------------------------------------------------- #

QDRANT_HOST: str = os.getenv("QDRANT_HOST", "localhost")
QDRANT_PORT: int = int(os.getenv("QDRANT_PORT", "6333"))
COLLECTION_NAME: str = os.getenv("QDRANT_COLLECTION", "rag_knowledge")

EMBEDDING_MODEL: str = os.getenv("OLLAMA_EMBEDDING_MODEL", "nomic-embed-text")
TEXT_SPLITTER_PARAMS = {
    "chunk_size": 1000,
    "chunk_overlap": 200,
}

# --------------------------------------------------------------------------- #
# Класс клиента Qdrant
# --------------------------------------------------------------------------- #

class QdrantClient:
    """
    Обёртка над qdrant_client и langchain_qdrant для упрощения работы с коллекцией.
    """

    def __init__(self, host: str = QDRANT_HOST, port: int = QDRANT_PORT,
                 collection_name: str = COLLECTION_NAME) -> None:
        self.host = host
        self.port = port
        self.collection_name = collection_name

        # Инициализируем низкоуровневый клиент
        self.client = _QdrantClient(host=self.host, port=self.port)

        # Создаём коллекцию при необходимости
        if not self._collection_exists():
            self._create_collection()

        # Обёртка LangChain для упрощённого добавления/поиска
        self.vector_store = QdrantVectorStore(
            client=self.client,
            collection_name=self.collection_name,
            embedding=OllamaEmbeddings(model=EMBEDDING_MODEL),
        )

    # ----------------------------------------------------------------------- #
    # Внутренние вспомогательные методы
    # ----------------------------------------------------------------------- #

    def _collection_exists(self) -> bool:
        """Проверяем наличие коллекции."""
        try:
            self.client.get_collection(name=self.collection_name)
            return True
        except Exception:
            return False

    def _create_collection(self) -> None:
        """Создаём коллекцию с параметрами по умолчанию."""
        vector_params = VectorParams(size=384, distance=Distance.COSINE)
        self.client.create_collection(
            collection_name=self.collection_name,
            vectors_config=vector_params
        )

    # ----------------------------------------------------------------------- #
    # Публичные методы
    # ----------------------------------------------------------------------- #

    def add_documents(self, documents: Iterable[str], titles: Iterable[str]) -> None:
        """
        Добавляет документы в коллекцию после разбиения на чанки.

        :param documents: Итерируемый список текстов.
        :param titles: Заголовки документов (один к одному с documents).
        """
        splitter = RecursiveCharacterTextSplitter(**TEXT_SPLITTER_PARAMS)
        all_chunks: List[Tuple[str, str]] = []

        for doc_text, title in zip(documents, titles):
            chunks = splitter.split_text(doc_text)
            # Для каждого чанка сохраняем заголовок как метаданные
            all_chunks.extend([(chunk, title) for chunk in chunks])

        texts, metas = zip(*all_chunks)

        self.vector_store.add_texts(
            texts=texts,
            metadatas=[{"title": meta} for meta in metas]
        )

    def search(self, query: str, limit: int = 5) -> List[dict]:
        """
        Семантический поиск по запросу.

        :param query: Текст запроса.
        :param limit: Количество результатов.
        :return: Список словарей с полями 'text', 'score' и метаданными.
        """
        results = self.vector_store.similarity_search_with_score(query, k=limit)
        return [
            {
                "text": r[0].page_content,
                "score": r[1],
                "metadata": r[0].metadata
            }
            for r in results
        ]

    # ----------------------------------------------------------------------- #
    # Утилиты для загрузки файлов из директории
    # ----------------------------------------------------------------------- #

    def load_from_directory(self, directory: str | Path) -> None:
        """
        Загружает все текстовые файлы из указанной папки в коллекцию.

        :param directory: Путь к каталогу.
        """
        dir_path = Path(directory)
        if not dir_path.is_dir():
            raise ValueError(f"Папка {directory} не существует")

        docs = []
        titles = []

        for file in dir_path.glob("**/*.txt"):
            text = file.read_text(encoding="utf-8")
            docs.append(text)
            titles.append(file.stem)

        self.add_documents(docs, titles)


# --------------------------------------------------------------------------- #
# Пример использования (можно удалить при импорте в другие модули)
# --------------------------------------------------------------------------- #

if __name__ == "__main__":
    client = QdrantClient()
    # Загрузка из папки ./data
    client.load_from_directory("./data")

    # Поиск по запросу
    query_text = "Какой смысл жизни?"
    results = client.search(query_text, limit=3)
    for i, res in enumerate(results, 1):
        print(f"\nРезультат {i}:")
        print(f"Score: {res['score']:.4f}")
        print(f"Title: {res['metadata'].get('title')}")
        print(f"Text snippet: {res['text'][:200]}...")