# qdrant_client.py
"""
Модуль инициализации клиента Qdrant, добавления документов и поиска по эмбеддингам Ollama.
Использует:
- langchain_qdrant.QdrantVectorStore для работы с коллекцией Qdrant
- langchain_ollama.OllamaEmbeddings для генерации векторных представлений
- RecursiveCharacterTextSplitter для разбиения больших текстов на чанки

Функциональность:
1. Инициализация клиента и создание/получение коллекции.
2. Добавление документа (с заголовком) в базу знаний с автоматическим чанкингом.
3. Поиск по семантическому запросу с ограничением количества результатов.

Author: ChatGPT
"""

from __future__ import annotations

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

from langchain_ollama import OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter


class QdrantClient:
    """
    Класс-обёртка над QdrantVectorStore.
    """

    def __init__(
        self,
        collection_name: str = "knowledge_base",
        host: str | None = None,
        port: int | None = None,
        embedding_model: str = "nomic-embed-text",
        chunk_size: int = 1000,
        chunk_overlap: int = 200,
    ) -> None:
        """
        Инициализация клиента Qdrant и подготовка коллекции.

        :param collection_name: имя коллекции в Qdrant
        :param host: хост Qdrant (по умолчанию localhost)
        :param port: порт Qdrant (по умолчанию 6333)
        :param embedding_model: название модели Ollama для эмбеддингов
        :param chunk_size: максимальная длина чанка в символах
        :param chunk_overlap: перекрытие между чанками
        """
        self.collection_name = collection_name

        # Параметры подключения к Qdrant
        host = host or os.getenv("QDRANT_HOST", "localhost")
        port = port or int(os.getenv("QDRANT_PORT", 6333))

        # Создаём объект эмбеддинговой модели Ollama
        self.embeddings = OllamaEmbeddings(model=embedding_model)

        # Инициализируем QdrantVectorStore
        self.store = QdrantVectorStore(
            url=f"http://{host}:{port}",
            collection_name=self.collection_name,
            embedding_function=self.embeddings.embed_query,  # используем embed_query для совместимости
        )

        # Создаём splitter для разбиения текста на чанки
        self.splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
        )

    def add_document(self, content: str, title: str) -> None:
        """
        Добавляет документ в базу знаний. Текст разбивается на чанки,
        каждому присваиваются метаданные (title и номер чанка).

        :param content: полный текст документа
        :param title: заголовок/имя документа
        """
        # Разбиваем контент на чанки
        chunks = self.splitter.split_text(content)

        documents = []
        for idx, chunk in enumerate(chunks):
            doc = {
                "page_content": chunk,
                "metadata": {"title": title, "chunk_index": idx},
            }
            documents.append(doc)

        # Добавляем векторные представления чанков
        self.store.add_documents(documents)

    def search(
        self,
        query: str,
        max_results: int = 5,
        filter_metadata: Dict[str, Any] | None = None,
    ) -> List[Dict[str, Any]]:
        """
        Семантический поиск по базе знаний.

        :param query: поисковый запрос
        :param max_results: максимальное количество результатов
        :param filter_metadata: словарь метаданных для фильтрации (например, {"title": "MyDoc"})
        :return: список найденных документов с полями:
                 - page_content
                 - metadata
                 - score
        """
        # Выполняем поиск векторного пространства
        results = self.store.similarity_search_with_score(
            query=query,
            k=max_results,
            filter=filter_metadata,
        )

        # Преобразуем к удобному формату
        formatted = []
        for doc, score in results:
            formatted.append(
                {
                    "content": doc.page_content,
                    "metadata": doc.metadata,
                    "score": score,
                }
            )
        return formatted


# --------------------------------------------------------------------------- #
# Пример использования (не входит в публичный API модуля)
# --------------------------------------------------------------------------- #

if __name__ == "__main__":
    # Инициализируем клиент
    client = QdrantClient()

    # Добавляем пример документа из файла
    sample_path = Path("sample.txt")
    if sample_path.exists():
        text = sample_path.read_text(encoding="utf-8")
        client.add_document(content=text, title=sample_path.stem)

    # Выполняем поиск
    query_str = "What is the main idea of the document?"
    results = client.search(query=query_str, max_results=3)
    for i, res in enumerate(results, 1):
        print(f"\nResult {i} (score={res['score']:.4f})")
        print(f"Title: {res['metadata'].get('title')}")
        print(f"Chunk index: {res['metadata'].get('chunk_index')}")
        print("Content snippet:", res["content"][:200], "...")