from dataclasses import dataclass
from typing import List

from tvalmetrics.llm_calls import ask_whether_answer_contains_context
from tvalmetrics.scorers.scorers_util import parse_boolean_response


@dataclass
class AugmentationAccuracy:
    """Augmentation accuracy score and information used to calculate the score.

    Fields
    ------
    score: float
        Float between 0 and 1 representing the augmentation accuracy score.
    answer: str
        The answer that was generated by the RAG system.
    context_list: List[str]
        Retrieved context used by the RAG system to make answer.
    answer_contains_context_list: List[bool]
        List of booleans representing whether information from each context in
        context_list is in answer.
    """

    score: float
    answer: str
    context_list: List[str]
    answer_contains_context_list: List[bool]


class AugmentationAccuracyScorer(object):
    """Class for calculating augmentation accuracy score.

    Parameters
    ----------
    model: str
        Name of the LLM model to use as the LLM evaluator.
    """

    def __init__(self, model: str):
        self.model = model

    def score(self, answer: str, context_list: List[str]) -> AugmentationAccuracy:
        """Calculate augmentation accuracy score.

        Parameters
        ----------
        answer: str
            The answer that was generated by the RAG system.
        context_list: List[str]
            Retrieved context used by the RAG system to make answer.

        Returns
        -------
        AugmentationAccuracyScore
            Augmentation accuracy score and information used to calculate the score.
        """
        contains_context_list = []
        for context in context_list:
            contains_context_response = ask_whether_answer_contains_context(
                answer, context, self.model
            )
            contains_context_list.append(
                parse_boolean_response(contains_context_response)
            )

        score = sum(contains_context_list) / len(contains_context_list)
        return AugmentationAccuracy(score, answer, context_list, contains_context_list)
