# -*- coding: utf-8 -*-
"""
Сервис переранжирования с использованием кросс-энкодеров.
Загружает модель один раз при первом вызове и использует её для ранжирования фрагментов.
"""

import logging
from typing import List, Tuple, Optional

logger = logging.getLogger(__name__)

# Импортируем sentence-transformers с обработкой ошибок
try:
    from sentence_transformers import CrossEncoder
    HAS_CROSS_ENCODER = True
except ImportError:
    HAS_CROSS_ENCODER = False
    logger.warning("sentence-transformers не установлен. Используйте LLM-переранжирование.")


class RerankerService:
    """
    Сервис переранжирования на основе кросс-энкодеров.
    Загружает модель при инициализации или при первом вызове.
    """

    # Рекомендуемая модель для русского языка
    DEFAULT_MODEL = "intfloat/multilingual-e5-reranker"

    def __init__(self, model_name: Optional[str] = None, device: Optional[str] = None):
        """
        Инициализация сервиса.

        Аргументы:
            model_name: имя модели (если None, используется DEFAULT_MODEL).
            device: устройство для вычислений ('cpu' или 'cuda'). Если None, автоопределение.
        """
        self.model_name = model_name or self.DEFAULT_MODEL
        self.device = device or ("cuda" if self._cuda_available() else "cpu")
        self._model = None
        logger.info(f"RerankerService инициализирован с моделью {self.model_name} (device={self.device})")

    @staticmethod
    def _cuda_available() -> bool:
        """Проверяет доступность CUDA (для GPU)."""
        try:
            import torch
            return torch.cuda.is_available()
        except ImportError:
            return False

    def _load_model(self) -> None:
        """Загружает модель кросс-энкодера (ленивая загрузка)."""
        if not HAS_CROSS_ENCODER:
            raise RuntimeError("sentence-transformers не установлен. Установите: pip install sentence-transformers")

        if self._model is None:
            logger.info(f"Загрузка модели {self.model_name} на устройство {self.device}...")
            self._model = CrossEncoder(self.model_name, device=self.device)
            logger.info(f"Модель {self.model_name} загружена")

    def rerank(
        self,
        query: str,
        fragments: List[str],
        top_k: Optional[int] = None
    ) -> List[Tuple[str, float]]:
        """
        Ранжирует фрагменты по релевантности запросу.

        Аргументы:
            query: текст запроса.
            fragments: список текстовых фрагментов.
            top_k: максимальное количество возвращаемых фрагментов (если None – все).

        Возвращает:
            List[Tuple[str, float]]: список кортежей (фрагмент, оценка_релевантности),
                                      отсортированный по убыванию оценки.
        """
        if not fragments:
            return []

        self._load_model()

        # Формируем пары (запрос, фрагмент) для оценки
        pairs = [(query, frag) for frag in fragments]

        # Получаем оценки релевантности (числа от 0 до 1)
        try:
            scores = self._model.predict(pairs)
        except Exception as e:
            logger.error(f"Ошибка при оценке кросс-энкодером: {e}")
            # В случае ошибки возвращаем фрагменты в исходном порядке с оценкой 0.0
            return [(frag, 0.0) for frag in fragments]

        # Собираем результаты и сортируем по убыванию оценки
        results = list(zip(fragments, scores))
        results.sort(key=lambda x: x[1], reverse=True)

        # Ограничиваем количество, если указано
        if top_k is not None:
            results = results[:top_k]

        logger.debug(f"Переранжирование: из {len(fragments)} фрагментов оставлено {len(results)}")
        return results

    def rerank_context(
        self,
        query: str,
        context: str,
        top_k: Optional[int] = None
    ) -> str:
        """
        Удобная обёртка: принимает строку контекста с фрагментами, разделёнными '\n\n',
        возвращает отранжированную строку.

        Аргументы:
            query: текст запроса.
            context: исходный контекст (фрагменты разделены '\n\n').
            top_k: максимальное количество возвращаемых фрагментов.

        Возвращает:
            str: отранжированный контекст (фрагменты разделены '\n\n').
        """
        if not context:
            return ""

        # Разбиваем контекст на фрагменты
        fragments = [f.strip() for f in context.split("\n\n") if f.strip()]
        if not fragments:
            return ""

        # Ранжируем
        ranked = self.rerank(query, fragments, top_k)

        # Собираем обратно в строку (только текст, без оценок)
        return "\n\n".join([frag for frag, _ in ranked])