from dataclasses import dataclass
from typing import List

from tvalmetrics.llm_calls import (
    ask_for_main_points,
    ask_whether_statement_derived_from_context,
)
from tvalmetrics.scorers.scorers_util import (
    parse_boolean_response,
    parse_bullet_list_response,
)


@dataclass
class AnswerConsistency:
    """Answer consistency score and information used to calculate the score.

    Fields
    ------
    score: float
        Float between 0 and 1 representing the answer consistency 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.
    main_point_list: List[str]
        List of main points of answer generated by the LLM evaluator.
    main_point_derived_from_context_list: List[bool]
        List of booleans representing whether each main point in main_point_list was
        derived from context in context_list.
    """

    score: float
    answer: str
    context_list: List[str]
    main_point_list: List[str]
    main_point_derived_from_context_list: List[bool]


class AnswerConsistencyScorer(object):
    """Class for calculating answer consistency 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]) -> AnswerConsistency:
        """Calculate answer consistency 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
        -------
        AnswerConsistencyScore
            The score and information used to calculate the score.
        """
        main_points_response = ask_for_main_points(answer, self.model)
        main_point_list = parse_bullet_list_response(main_points_response)
        main_point_derived_from_context_list = []
        for main_point in main_point_list:
            statement_derived_from_context_response = (
                ask_whether_statement_derived_from_context(
                    main_point, context_list, self.model
                )
            )
            main_point_derived_from_context_list.append(
                parse_boolean_response(statement_derived_from_context_response)
            )
        score = sum(main_point_derived_from_context_list) / len(main_point_list)
        return AnswerConsistency(
            score,
            answer,
            context_list,
            main_point_list,
            main_point_derived_from_context_list,
        )
