"""
Contextual Recall Test
"""

import json

from .prompt_template import ContextualRecallTemplate


class ContextualRecallTest:
    """
    A class for testing contextual recall.
    """

    def __init__(
        self,
        client,
        expected_output,
        retrieval_context,
        include_reason=True,
        model="gpt-3.5-turbo",
        threshold=0.5,
        temperature=0,
    ):
        """
        Initializes the ContextualRecallTest instance with the given parameters

        Args:
        expected_output (str): The expected output for the test
        retrieval_context (list): A list of strings representing the retrieval context
        include_reason (bool): Whether to include the reason in the test result (default True)
        model (str): The name of the OpenAI model to use (default "gpt-3.5-turbo")
        threshold (float): The threshold for the test score (default 0.5)
        temperature (int): The temperature parameter for the model (default 0)
        """
        self.expected_output = expected_output
        self.retrieval_context = retrieval_context
        self.include_reason = include_reason
        self.model_name = model
        self.threshold = threshold
        self.temperature = temperature
        self.client = client

    def model(self, prompt):
        """
        Generates completions for the given prompt using the specified model

        Args:
        prompt (str): The prompt for generating completions

        Returns:
        dict: The completions generated by the model
        """
        return self.client.chat.completions.create(
            model=self.model_name,
            temperature=self.temperature,
            messages=[{"role": "system", "content": prompt}],
        )

    def trim_to_json(self, response):
        """
        Trims the response to extract the JSON content

        Args:
        response (dict): The response from the model

        Returns:
        dict: The JSON content extracted from the response
        """
        message_content = response.choices[0].message.content
        data = json.dumps(message_content)
        return json.loads(data)

    def generate_verdicts(self):
        """
        Generates verdicts based on the expected output and retrieval context

        Returns:
        list: A list of generated verdicts
        """
        verdict_prompt = ContextualRecallTemplate.generate_verdicts(
            self.expected_output, self.retrieval_context
        )
        verdict_response = self.model(verdict_prompt)
        verdicts = json.loads(self.trim_to_json(verdict_response))["verdicts"]
        return verdicts

    def generate_score(self, verdicts):
        """
        Generates a score based on the generated verdicts

        Args:
        verdicts (list): A list of generated verdicts

        Returns:
        float: The calculated score
        """
        if len(verdicts) == 0:
            return 0

        justified_sentences = sum(
            1 for verdict in verdicts if verdict["verdict"].lower() == "yes"
        )
        return justified_sentences / len(verdicts)

    def generate_reason(self, supportive_reasons, unsupportive_reasons, score):
        """
        Generates a reason based on the supportive and unsupportive reasons

        Args:
        supportive_reasons (list): A list of supportive reasons
        unsupportive_reasons (list): A list of unsupportive reasons
        score (float): The calculated score

        Returns:
        dict: The generated reason
        """
        prompt = ContextualRecallTemplate.generate_reason(
            expected_output=self.expected_output,
            supportive_reasons=supportive_reasons,
            unsupportive_reasons=unsupportive_reasons,
            score=score,
        )

        res = self.model(prompt)
        reason = self.trim_to_json(res)
        return reason

    def run(self):
        """
        Runs the test and generates the result

        Returns:
        dict: The test result
        """
        verdicts = self.generate_verdicts()
        score = self.generate_score(verdicts)

        supportive_reasons = [
            verdict["reason"]
            for verdict in verdicts
            if verdict["verdict"].lower() == "yes"
        ]
        unsupportive_reasons = [
            verdict["reason"]
            for verdict in verdicts
            if verdict["verdict"].lower() != "yes"
        ]

        result = {
            "expected_response": self.expected_output,
            "context": self.retrieval_context,
            "score": score,
            "threshold": self.threshold,
            "is_passed": score >= self.threshold,
            "evaluated_with": {
                "model": self.model_name,
                "include_reason": self.include_reason,
            },
        }

        if self.include_reason:
            reason = self.generate_reason(
                supportive_reasons, unsupportive_reasons, score
            )
            result["reason"] = reason

        return result


# if __name__ == "__main__":
#     expected_response = "Paris"
#     context = ["Paris is the capital of France.", "london is a city"]
#     test_instance = ContextualRecallTest(expected_response, context)
#     result = test_instance.run()
#     print(result)
