import logging
import json
import os
from typing import Dict, Any, List
from benchmarks.base import Benchmark
from constants import TEST_SEPARATOR

class SummarizationBenchmark(Benchmark):
    """Бенчмарк для тестирования пересказов."""

    def __init__(self):
        super().__init__("summarization")
        # Загружаем универсальный промпт
        with open('prompts/summarization.txt', 'r', encoding='utf-8') as f:
            self.universal_prompt = f.read().strip()

    def load_test_data(self) -> List[Dict[str, Any]]:
        """
        Загрузка тестовых данных для пересказов.

        Returns:
            Список тестовых случаев
        """
        test_data = []
        data_dir = "tests/summarization"

        for filename in os.listdir(data_dir):
            if filename.endswith('.txt'):
                with open(os.path.join(data_dir, filename), 'r', encoding='utf-8') as f:
                    content = f.read()
                    # Разделяем по разделителю
                    parts = content.split(TEST_SEPARATOR, 1)
                    if len(parts) == 2:
                        test_data.append({
                            'name': filename.replace('.txt', ''),
                            'prompt': self.universal_prompt.format(text=parts[0]),
                            'expected': parts[1]
                        })

        return test_data

    def evaluate(self, model_response: str, expected: str) -> float:
        """
        Оценка качества пересказа.

        Args:
            model_response: Ответ от модели
            expected: Ожидаемый ответ

        Returns:
            Метрика качества (0-1)
        """
        # Простая оценка на основе совпадения токенов
        model_tokens = set(model_response.lower().split())
        expected_tokens = set(expected.lower().split())

        if len(expected_tokens) == 0:
            return 0.0

        intersection = model_tokens.intersection(expected_tokens)
        precision = len(intersection) / len(model_tokens) if model_tokens else 0.0
        recall = len(intersection) / len(expected_tokens) if expected_tokens else 0.0

        # F1-score
        if (precision + recall) == 0:
            return 0.0
        f1 = 2 * (precision * recall) / (precision + recall)

        return round(f1, 3)
