from dataclasses import dataclass
from typing import List

from tvalmetrics.scorers.augmentation_accuracy_scorer import AugmentationAccuracyScorer
from tvalmetrics.scorers.augmentation_precision_scorer import (
    AugmentationPrecisionScorer,
)
from tvalmetrics.scorers.retrieval_precision_scorer import RetrievalPrecisionScorer


@dataclass
class ContextScores:
    """Retrieval precision, augmentation precision, and augmentation accuracy scores.

    Fields
    ------
    retrieval_precision: float
        Float between 0 and 1 representing the retrieval precision score.
    augmentation_precision: float
        Float between 0 and 1 representing the augmentation precision score.
    augmentation_accuracy: float
        Float between 0 and 1 representing the augmentation accuracy score."""

    retrieval_precision: float
    augmentation_precision: float
    augmentation_accuracy: float


class ContextScorer(object):
    """Class for calculating the three main context scores at the same time.

    Calculates retrieval precision, augmentation precision, and augmentation accuracy at
    the same time to avoid repeating the same LLM calls that would happen if each score
    was calculated separately.

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

    def __init__(self, model: str):
        self.model = model
        self.retrieval_precision_scorer = RetrievalPrecisionScorer(self.model)
        self.augmentation_accuracy_scorer = AugmentationAccuracyScorer(self.model)

    def score(
        self, question: str, answer: str, context_list: List[str]
    ) -> ContextScores:
        """Calculate the three main context scores.

        Caclulates retrieval precision, augmentation precision, and augmentation
        accuracy.

        Parameters
        ----------
        question: str
            The question that was asked.
        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
        -------
        ContextScores
            Retrieval precision, augmentation precision, and augmentation accuracy
            scores.
        """
        retrieval_precision_score = self.retrieval_precision_scorer.score(
            question, context_list
        )
        retrieval_precision = retrieval_precision_score.score
        context_relevant_list = retrieval_precision_score.context_relevant_list
        augmentation_accuracy_score = self.augmentation_accuracy_scorer.score(
            answer, context_list
        )
        augmentation_accuracy = augmentation_accuracy_score.score
        contains_context_list = augmentation_accuracy_score.answer_contains_context_list
        augmentation_precision = AugmentationPrecisionScorer.score_from_context_labels(
            context_relevant_list, contains_context_list
        )

        return ContextScores(
            retrieval_precision=retrieval_precision,
            augmentation_precision=augmentation_precision,
            augmentation_accuracy=augmentation_accuracy,
        )
