"""Клиент для работы с векторной базой данных Milvus."""
import logging
import os
from typing import Dict, List, Optional

from pymilvus import Collection, SearchResult, connections

from data.processors import filter_short_and_invalid_words

logger = logging.getLogger(__name__)


class MilvusConnectionManager:
    """Управляет подключениями к векторной базе данных Milvus."""

    def __init__(self) -> None:
        """Инициализирует менеджер подключений с переменными окружения."""
        self.host = os.getenv('MILVUS_HOST')
        self.port = os.getenv('MILVUS_PORT')
        self.database_name = os.getenv('DATABASE_NAME')
        self.connection_active = False

    def connect(self) -> None:
        """Устанавливает соединение с Milvus, если оно еще не активно."""
        if not self.connection_active:
            connections.connect(host=self.host, port=self.port, db_name=self.database_name)
            self.connection_active = True
            logger.info('Connected to Milvus database: %s', self.database_name)

    def disconnect(self) -> None:
        """Отключается от Milvus, если соединение активно."""
        if self.connection_active:
            connections.disconnect(alias=self.database_name)
            self.connection_active = False
            logger.info('Disconnected from Milvus database: %s', self.database_name)

    def get_collection(self, collection_name: str) -> Optional[Collection]:
        """Получает коллекцию и загружает её в память.

        Args:
            collection_name: Название коллекции

        Returns:
            Объект коллекции Milvus или None при ошибке
        """
        try:
            collection = Collection(name=collection_name)
            collection.load()
            return collection
        except Exception as error:
            logger.error('Failed to get collection %s: %s', collection_name, error)
            return None

    def release_collection(self, collection: Collection) -> None:
        """Освобождает ресурсы коллекции."""
        if collection is not None:
            collection.release()


class MilvusTableSearch:
    """Обеспечивает поиск информации о таблицах в Milvus."""

    def __init__(self, milvus_manager: MilvusConnectionManager) -> None:
        """Инициализирует с менеджером подключений Milvus.

        (MilvusConnectionManager)
        """
        self.milvus_manager = milvus_manager

    def search_by_bm25_embeddings(
        self,
        collection: Collection,
        text: str,
        limit: int,
    ) -> SearchResult:

        search_results = collection.search(
            data=[text],
            limit=limit,
            anns_field='embedding_bm25',
            param={"params": {"drop_ratio_search": 0.0}},
            output_fields=['table_name', 'table_info'],
        )

        return search_results

    def filter_non_zero_distance_results(
        self,
        search_results: SearchResult,
        num_table_out: int,
    ) -> List[Dict[str, str]]:
        """Фильтрует результаты, оставляя только те, у которых distance > 0.

        Отбрасывает неподходящие результаты и ограничивает количество возвращаемых записей.
        Args:
            search_results: Результаты поиска в формате SearchResult от Milvus
        """
        filtered_results: List[Dict[str, str]] = []

        valid_hits = (hit for hits in search_results for hit in hits if hit.distance > 0)

        for hit in valid_hits:
            filtered_results.append(
                {
                    'table_name': hit.entity.get('table_name'),
                    'table_info': hit.entity.get('table_info'),
                    'distance': hit.distance,
                },
            )

            if len(filtered_results) >= num_table_out:
                break

        return filtered_results

    def search_table_info(self, user_question: str, reference_project: str, num_table_out: int = 4) -> str:
        """Поиск информации о таблицах в Milvus на основе запроса пользователя.

        Args:
            user_question: Запрос пользователя
            reference_project: Название проекта для контекста
            num_table_out: Количество таблиц для возврата

        Returns:
            Строка с информацией о найденных таблицах или пустая строка, если ничего не найдено
        """

        table_info_results = []
        COLLECTION_NAME = f'xml_{reference_project}'
        collection = self.milvus_manager.get_collection(COLLECTION_NAME)

        if not collection:
            return ''

        logger.info('Searching by BM25 embeddings...')
        try:
            bm25_results = self.search_by_bm25_embeddings(
                collection,
                filter_short_and_invalid_words(user_question),
                num_table_out,
            )

        except Exception as error:
            logger.error('Error searching in Milvus table: %s', str(error))
            return ''
        finally:
            if collection:
                self.milvus_manager.release_collection(collection)

        filtered_results = self.filter_non_zero_distance_results(bm25_results, num_table_out)

        for cur_result in filtered_results:
            table_info_results.append(f"Таблица: {cur_result['table_name']}, Инфо: {cur_result['table_info']}")

        if not filtered_results:
            logger.info('No results found with distance > 0.')

        if table_info_results:
            joined_tables = '\n'.join(table_info_results)
            return f'Информация о таблицах:\n{joined_tables}'
        else:
            return ''
