from braintrust_core.score import Score

from autoevals.llm import OpenAIScorer

from .oai import arun_cached_request, run_cached_request

REQUEST_TYPE = "moderation"


class Moderation(OpenAIScorer):
    """
    A scorer that uses OpenAI's moderation API to determine if AI response contains ANY flagged content.
    """

    threshold = None
    extra_args = {}

    def __init__(self, threshold=None, api_key=None, base_url=None):
        """
        Create a new Moderation scorer.

        :param threshold: Optional. Threshold to use to determine whether content has exceeded threshold. By
        default, it uses OpenAI's default. (Using `flagged` from the response payload.)
        :param api_key: OpenAI key
        :param base_url: Base URL to be used to reach OpenAI moderation endpoint.
        """
        super().__init__(api_key=api_key, base_url=base_url)
        self.threshold = threshold

    def _run_eval_sync(self, output, __expected=None):
        moderation_response = run_cached_request(REQUEST_TYPE, input=output, **self.extra_args)["results"][0]
        return self.__postprocess_response(moderation_response)

    def __postprocess_response(self, moderation_response) -> Score:
        return Score(
            name=self._name(),
            score=self.compute_score(moderation_response, self.threshold),
            metadata={
                "threshold": self.threshold,
                "category_scores": moderation_response["category_scores"],
            },
        )

    async def _run_eval_async(self, output, expected=None, **kwargs) -> Score:
        moderation_response = (await arun_cached_request(REQUEST_TYPE, input=output, **self.extra_args))["results"][0]
        return self.__postprocess_response(moderation_response)

    @staticmethod
    def compute_score(moderation_result, threshold):
        if threshold is None:
            return 0 if moderation_result["flagged"] else 1

        category_scores = moderation_result["category_scores"]
        for category in category_scores.keys():
            if category_scores[category] > threshold:
                return 0

        return 1


__all__ = ["Moderation"]
