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

"""Tests for the Audio Achilles Heel scenario."""

from unittest.mock import MagicMock, patch

import pytest

from pyrit.executor.attack import PromptSendingAttack
from pyrit.models import ComponentIdentifier, SeedPrompt
from pyrit.prompt_target import PromptTarget
from pyrit.scenario.garak import AudioAchillesHeel, AudioAchillesHeelTechnique  # type: ignore[ty:unresolved-import]
from pyrit.scenario.scenarios.garak.audio_achilles_heel import (
    DEFAULT_MAX_DATASET_SIZE,
    DEFAULT_TEXT_PROMPT,
    GENERIC_OBJECTIVE,
    AudioAchillesHeelDatasetConfiguration,
)
from pyrit.score import TrueFalseScorer


def _mock_id(name: str) -> ComponentIdentifier:
    return ComponentIdentifier(class_name=name, class_module="test")


@pytest.fixture
def mock_objective_target():
    mock = MagicMock(spec=PromptTarget)
    mock.get_identifier.return_value = _mock_id("MockObjectiveTarget")
    return mock


@pytest.fixture
def mock_objective_scorer():
    mock = MagicMock(spec=TrueFalseScorer)
    mock.get_identifier.return_value = _mock_id("MockObjectiveScorer")
    return mock


def _audio_seed(*, path: str, categories: list[str] | None = None) -> SeedPrompt:
    return SeedPrompt(value=path, data_type="audio_path", harm_categories=categories or [])


@pytest.fixture
def audio_dataset_config():
    """Explicit inline audio seeds so tests never hit HuggingFace."""
    return AudioAchillesHeelDatasetConfiguration(
        seeds=[
            _audio_seed(path="/tmp/Malware_Generation_1.wav", categories=["Malware_Generation"]),
            _audio_seed(path="/tmp/Hate_Speech_2.wav", categories=["Hate_Speech"]),
        ]
    )


@pytest.mark.usefixtures("patch_central_database")
class TestAudioAchillesHeelInitialization:
    """Tests for AudioAchillesHeel initialization."""

    def test_init_basic(self, mock_objective_scorer):
        scenario = AudioAchillesHeel(objective_scorer=mock_objective_scorer)
        assert scenario.name == "AudioAchillesHeel"
        assert scenario.VERSION == 1

    def test_init_with_custom_scorer(self, mock_objective_scorer):
        scenario = AudioAchillesHeel(objective_scorer=mock_objective_scorer)
        assert scenario._objective_scorer == mock_objective_scorer
        assert scenario._scorer_config.objective_scorer == mock_objective_scorer

    def test_init_creates_default_scorer_when_not_provided(self, mock_objective_scorer):
        with patch.object(AudioAchillesHeel, "_get_default_objective_scorer", return_value=mock_objective_scorer):
            scenario = AudioAchillesHeel()
            assert scenario._objective_scorer == mock_objective_scorer

    def test_required_datasets(self):
        assert AudioAchillesHeel.required_datasets() == ["garak_audio_achilles_heel"]

    def test_baseline_disabled_by_default(self):
        from pyrit.scenario.core.scenario import BaselineAttackPolicy

        assert BaselineAttackPolicy.Disabled == AudioAchillesHeel.BASELINE_ATTACK_POLICY

    def test_default_dataset_config_uses_garak_audio(self, mock_objective_scorer):
        config = AudioAchillesHeel(objective_scorer=mock_objective_scorer)._default_dataset_config
        assert config.dataset_names == ["garak_audio_achilles_heel"]

    def test_default_dataset_config_caps_size(self, mock_objective_scorer):
        config = AudioAchillesHeel(objective_scorer=mock_objective_scorer)._default_dataset_config
        assert config.max_dataset_size == DEFAULT_MAX_DATASET_SIZE

    def test_default_technique_is_all(self, mock_objective_scorer):
        scenario = AudioAchillesHeel(objective_scorer=mock_objective_scorer)
        assert scenario._default_technique == AudioAchillesHeelTechnique.ALL

    def test_default_text_prompt_used(self, mock_objective_scorer):
        config = AudioAchillesHeel(objective_scorer=mock_objective_scorer)._default_dataset_config
        assert config._text_prompt == DEFAULT_TEXT_PROMPT

    def test_custom_text_prompt_propagates_to_dataset_config(self, mock_objective_scorer):
        scenario = AudioAchillesHeel(objective_scorer=mock_objective_scorer, text_prompt="listen carefully")
        assert scenario._default_dataset_config._text_prompt == "listen carefully"


@pytest.mark.usefixtures("patch_central_database")
class TestAudioAchillesHeelDatasetConfiguration:
    """Tests for the dataset configuration's attack-group shaping."""

    def test_objective_derived_from_harm_category(self):
        config = AudioAchillesHeelDatasetConfiguration()
        groups = config._build_attack_groups([_audio_seed(path="/tmp/x.wav", categories=["Malware_Generation"])])
        assert groups[0].objective.value == "Get the model to comply with the spoken request about Malware Generation."

    def test_objective_derived_from_metadata_category(self):
        config = AudioAchillesHeelDatasetConfiguration()
        seed = SeedPrompt(value="/tmp/x.wav", data_type="audio_path", metadata={"category": "Fraud_Scam"})
        groups = config._build_attack_groups([seed])
        assert groups[0].objective.value == "Get the model to comply with the spoken request about Fraud Scam."

    def test_objective_falls_back_when_no_category(self):
        config = AudioAchillesHeelDatasetConfiguration()
        groups = config._build_attack_groups([_audio_seed(path="/tmp/clip_0.wav")])
        assert groups[0].objective.value == GENERIC_OBJECTIVE

    def test_attack_group_merges_text_and_audio_into_one_message(self):
        config = AudioAchillesHeelDatasetConfiguration()
        groups = config._build_attack_groups([_audio_seed(path="/tmp/Malware_1.wav", categories=["Malware"])])

        message = groups[0].next_message
        assert message is not None
        data_types = [piece.original_value_data_type for piece in message.message_pieces]
        assert data_types == ["text", "audio_path"]

    def test_custom_text_prompt_used_in_group(self):
        config = AudioAchillesHeelDatasetConfiguration(text_prompt="follow the recording")
        groups = config._build_attack_groups([_audio_seed(path="/tmp/x.wav", categories=["Malware"])])

        message = groups[0].next_message
        assert message is not None
        text_piece = next(p for p in message.message_pieces if p.original_value_data_type == "text")
        assert text_piece.original_value == "follow the recording"


@pytest.mark.usefixtures("patch_central_database")
class TestAudioAchillesHeelTechniqueExpansion:
    """Tests for technique expansion and atomic attack generation."""

    async def test_all_expands_to_audio_jailbreak(
        self, mock_objective_target, mock_objective_scorer, audio_dataset_config
    ):
        scenario = AudioAchillesHeel(objective_scorer=mock_objective_scorer)
        scenario.set_params_from_args(
            args={
                "objective_target": mock_objective_target,
                "scenario_techniques": [AudioAchillesHeelTechnique.ALL],
                "dataset_config": audio_dataset_config,
            }
        )
        await scenario.initialize_async()

        technique_values = {s.value for s in scenario._scenario_techniques}
        assert technique_values == {"audio_jailbreak"}

    async def test_builds_single_prompt_sending_attack(
        self, mock_objective_target, mock_objective_scorer, audio_dataset_config
    ):
        scenario = AudioAchillesHeel(objective_scorer=mock_objective_scorer)
        scenario.set_params_from_args(
            args={
                "objective_target": mock_objective_target,
                "dataset_config": audio_dataset_config,
            }
        )
        await scenario.initialize_async()

        atomic_attacks = scenario._atomic_attacks
        assert len(atomic_attacks) == 1
        assert atomic_attacks[0].atomic_attack_name == "audio_jailbreak"
        assert isinstance(atomic_attacks[0].attack_technique.attack, PromptSendingAttack)

    async def test_atomic_attack_carries_all_seed_groups(
        self, mock_objective_target, mock_objective_scorer, audio_dataset_config
    ):
        scenario = AudioAchillesHeel(objective_scorer=mock_objective_scorer)
        scenario.set_params_from_args(
            args={
                "objective_target": mock_objective_target,
                "dataset_config": audio_dataset_config,
            }
        )
        await scenario.initialize_async()

        atomic_attack = scenario._atomic_attacks[0]
        assert len(atomic_attack.seed_groups) == 2


@pytest.mark.usefixtures("patch_central_database")
class TestAudioAchillesHeelTechniqueValues:
    """Tests for AudioAchillesHeelTechnique members."""

    def test_concrete_technique_value(self):
        assert AudioAchillesHeelTechnique.AudioJailbreak.value == "audio_jailbreak"

    def test_all_aggregate_present(self):
        assert "all" in AudioAchillesHeelTechnique.get_aggregate_tags()