"""Сервис для переранжирования документов с использованием кросс-энкодера."""
import logging
import os
from typing import List, Tuple

import requests

logger = logging.getLogger(__name__)


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

    def __init__(self) -> None:
        """Инициализирует сервис переранжирования."""
        self.url = os.getenv('URL_RERANKER')
        if not self.url:
            logger.warning('URL_RERANKER environment variable not set')

    def rank_documents(
        self,
        query: str,
        documents: List[str],
        document_page_infos: List[Tuple[str]],
        model_name: str,
        truncate_prompt_tokens: int = 512,
        top_k: int = 10,
    ) -> List[Tuple[str, Tuple[str], float]]:
        """Ранжирует список документов с использованием кросс-энкодера.
        
        Отправляет запрос к API переранжирования, которое оценивает релевантность каждого 
        документа к запросу. Возвращает отсортированный по релевантности список документов.
        
        Args:
            query: Текст запроса пользователя
            documents: Список текстов документов для ранжирования
            document_page_infos: Список кортежей с page_link и page_title для каждого документа
            model_name: Название модели кросс-энкодера
            truncate_prompt_tokens: Максимальное количество токенов для усечения промпта
            top_k: Количество возвращаемых топ-документов
            
        Returns:
            Список кортежей (документ, информация о странице, оценка) в порядке убывания релевантности.
            В случае ошибки возвращает top_k документов с нулевыми оценками.
        """
        zero_scores = [0.0 for _ in range(len(documents))]
        if not self.url:
            logger.error('URL_RERANKER not set, cannot rerank documents')
            return list(zip(documents, document_page_infos, zero_scores))[:top_k]

        json_data = {
            'text_1': [query],
            'text_2': documents,
            'model': model_name,
            'truncate_prompt_tokens': truncate_prompt_tokens,
        }

        try:
            response = requests.post(self.url, headers={'Content-Type': 'application/json'}, json=json_data)

            if response.status_code == requests.codes.ok:
                scores = response.json()['data']
                ranked_items = sorted(
                    zip(documents, document_page_infos, scores),
                    key=lambda ranked_item: ranked_item[2]['score'],
                    reverse=True,
                )[:top_k]

                return [(doc, doc_id, scores['score']) for doc, doc_id, scores in ranked_items]
            else:
                logger.error('Cross-encoder error: %s, %s', response.status_code, response.text)
                return list(zip(documents, document_page_infos, zero_scores))[:top_k]

        except Exception as error:
            logger.error('Error during document reranking: %s', str(error))
            return list(zip(documents, document_page_infos, zero_scores))[:top_k]
