# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import json
import logging
import uuid
from typing import Any

from pyrit.models import ComponentIdentifier, Message, MessagePiece, Score, ScoreType
from pyrit.prompt_target import PromptShieldTarget
from pyrit.score.scorer_prompt_validator import ScorerPromptValidator
from pyrit.score.true_false.true_false_score_aggregator import (
    TrueFalseAggregatorFunc,
    TrueFalseScoreAggregator,
)
from pyrit.score.true_false.true_false_scorer import MessageTrueFalseScorer

logger = logging.getLogger(__name__)


class PromptShieldScorer(MessageTrueFalseScorer):
    """
    Returns true if an attack or jailbreak has been detected by Prompt Shield.
    """

    scorer_type: ScoreType
    _prompt_shield_target: PromptShieldTarget

    _DEFAULT_VALIDATOR: ScorerPromptValidator = ScorerPromptValidator(supported_data_types=["text"])

    def __init__(
        self,
        *,
        prompt_shield_target: PromptShieldTarget,
        validator: ScorerPromptValidator | None = None,
        score_aggregator: TrueFalseAggregatorFunc = TrueFalseScoreAggregator.OR,
    ) -> None:
        """
        Initialize the PromptShieldScorer.

        Args:
            prompt_shield_target (PromptShieldTarget): The Prompt Shield target to use for scoring.
            validator (ScorerPromptValidator | None): Custom validator. Defaults to None.
            score_aggregator (TrueFalseAggregatorFunc): The aggregator function to use.
                Defaults to TrueFalseScoreAggregator.OR.
        """
        self._prompt_target = prompt_shield_target

        super().__init__(validator=validator or self._DEFAULT_VALIDATOR, score_aggregator=score_aggregator)

    def _build_identifier(self) -> ComponentIdentifier:
        """
        Build the identifier for this scorer.

        Returns:
            ComponentIdentifier: The identifier for this scorer.
        """
        return self._create_identifier(
            score_aggregator=self._score_aggregator.__name__,  # type: ignore[ty:unresolved-attribute]
            prompt_target=self._prompt_target.get_identifier(),
        )

    async def _score_piece_async(self, message_piece: MessagePiece, *, objective: str | None = None) -> list[Score]:
        conversation_id = str(uuid.uuid4())

        body = message_piece.original_value

        request = Message(
            message_pieces=[
                MessagePiece(
                    role="user",
                    original_value=body,
                    prompt_metadata=message_piece.prompt_metadata,
                    conversation_id=conversation_id,
                )
            ]
        )

        # The body of the Prompt Shield response
        target_response = await self._prompt_target.send_prompt_async(message=request)
        response: str = target_response[0].get_value()

        # Whether or not any of the documents or userPrompt got flagged as an attack
        result: bool = any(self._parse_response_to_boolean_list(response))

        # Store parsed response in metadata for structured access
        try:
            meta = json.loads(response)
        except Exception:
            meta = {"raw": response}

        score = Score(
            score_type="true_false",
            score_value=str(result),
            score_value_description="True if an attack or jailbreak has been detected, else False.",
            score_category=["attack_detection"],
            score_metadata=meta,  # type: ignore[ty:invalid-argument-type]
            score_rationale="",
            scorer_class_identifier=self.get_identifier(),
            message_piece_id=message_piece.id,
            objective=objective,
        )

        return [score]

    def _parse_response_to_boolean_list(self, response: str) -> list[bool]:
        """
        Remember that you can just access the metadata attribute to get the original Prompt Shield endpoint response,
        and then just call json.loads() on it to interact with it.

        Returns:
            list[bool]: A list of boolean values indicating whether an attack was detected.
        """
        response_json: dict[str, Any] = json.loads(response)

        user_prompt_attack: dict[str, bool] = response_json.get("userPromptAnalysis", False)
        documents_attack: list[dict[str, Any]] = response_json.get("documentsAnalysis", False)

        user_detections: list[bool] = (
            [False] if not user_prompt_attack else [bool(user_prompt_attack.get("attackDetected"))]
        )

        if not documents_attack:
            document_detections: list[bool] = [False]
        else:
            document_detections = [bool(document.get("attackDetected")) for document in documents_attack]

        return user_detections + document_detections