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

import asyncio
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
        self._load_lock = asyncio.Lock()
        self._model_loaded = asyncio.Event()
        self._loading = False
        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

    async def _load_model_async(self) -> None:
        """Загружает модель кросс-энкодера асинхронно (ленивая загрузка)."""
        async with self._load_lock:
            if self._model_loaded.is_set():
                return
            if self._loading:
                await self._model_loaded.wait()
                return
            self._loading = True
            try:
                loop = asyncio.get_running_loop()
                self._model = await loop.run_in_executor(None, self._load_model_sync)
                self._model_loaded.set()
                logger.info(f"Модель {self.model_name} загружена")
            except Exception as e:
                logger.error(f"Ошибка загрузки модели: {e}")
                self._loading = False
                raise
            finally:
                self._loading = False

    def _load_model_sync(self):
        """Синхронная загрузка модели (выполняется в отдельном потоке)."""
        if not HAS_CROSS_ENCODER:
            raise RuntimeError("sentence-transformers не установлен. Установите: pip install sentence-transformers")
        return CrossEncoder(self.model_name, device=self.device)

    async 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 []

        await self._load_model_async()

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

        try:
            # Выполняем предсказание синхронно в отдельном потоке
            loop = asyncio.get_running_loop()
            scores = await loop.run_in_executor(None, self._model.predict, pairs)
            results = list(zip(fragments, scores))
            results.sort(key=lambda x: x[1], reverse=True)
            if top_k:
                results = results[:top_k]
            return results
        except Exception as e:
            logger.error(f"Ошибка при оценке кросс-энкодером: {e}")
            return [(frag, 0.0) for frag in fragments]

    async 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 = await self.rerank(query, fragments, top_k)
        return "\n\n".join([frag for frag, _ in ranked])