"""Модуль генерации текста для пайплайна."""
import logging
import random
import time
from typing import Dict, List, Optional, Tuple

from data.processors import create_few_shot_prompt
from models.llm import LLMService
from utils.metrics import calculate_average

logger = logging.getLogger(__name__)


class TextGenerator:
    """Сервис для генерации текста на основе контекста и примеров с
    использованием языковых моделей."""

    def __init__(self, llm_service: LLMService) -> None:
        """Инициализация генератора текста (LLMService)"""
        self.llm_service = llm_service

    def collect_previous_generations(
        self,
        all_answers: List[str],
        model_name: str,
        is_summary: Optional[bool] = False,
    ) -> str:
        """Собирает и суммаризирует предыдущие генерации текста.

        Returns:
            - Суммаризованный текст или пустая строка
        """
        if not is_summary or not all_answers:
            return ''

        previous_generations = []
        for answer in all_answers:
            summary = self.llm_service.summarize(answer, model_name)
            previous_generations.append(summary)

        return '\n'.join(previous_generations)

    def generate(
        self,
        user_question: str,
        template: str,
        context: str,
        system_message: str,
        examples: List[Dict[str, str]],
        model_name: str,
        num_generations: int,
        reference_answer: str,
        previous_generations: str = ''
    ) -> Tuple[str, Dict[str, float], float, Dict[str, float]]:
        """Генерирует текст на основе контекста и примеров.

        Создает промпт в формате few-shot, генерирует несколько вариантов ответов,
        рассчитывает метрики и выбирает случайный ответ из сгенерированных.
        При наличии логгера записывает результаты, метрики и статистику токенов в postgres

        Args:
            user_question: Вопрос пользователя
            template: Шаблон для генерации ответа
            context: Контекстная информация из базы знаний
            system_message: Системное сообщение для LLM
            examples: Примеры для few-shot обучения
            model_name: Имя модели для генерации
            num_generations: Количество вариантов генерации
            reference_answer: Эталонный ответ для расчета метрик
            previous_generations: Суммаризированные предыдущие генерации

        Returns:
            - сгенерированный_ответ
            - метрики
            - время_генерации
            - статистика_токенов
        """
        messages = create_few_shot_prompt(
            system_message=system_message,
            examples=examples,
            user_question=user_question,
            template=template,
            context=context,
            previous_generations=previous_generations,
        )
        
        start = time.perf_counter()
        metrics_list, generated_answers, tokens_list = self.llm_service.generate_and_validate(
            messages=messages,
            model_name=model_name,
            reference_answer=reference_answer,
            user_input=user_question + template,
            documents=[context],
            num_generations=num_generations,
        )
        generation_time = (time.perf_counter() - start) / num_generations

        generated_answer = random.choice(generated_answers)

        avg_metrics = calculate_average(metrics_list)
        avg_tokens = calculate_average(tokens_list)

        return generated_answer, avg_metrics, generation_time, avg_tokens
