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

from unittest.mock import patch

import pytest

from pyrit.datasets.seed_datasets.remote.mlcommons_ailuminate_dataset import _MLCommonsAILuminateDataset
from pyrit.models import SeedDataset, SeedPrompt

AILUMINATE_EXPECTED_HAZARD_CATEGORIES = [
    ("vcr", ["VIOLENT_CONTENT"]),
    ("src", ["SEXUAL_CONTENT"]),
    ("cse", ["SEXUAL_CONTENT", "GROOMING"]),
    ("ssh", ["SUICIDE", "SELF_HARM"]),
    ("iwp", ["CBRN", "MILITARY"]),
    ("ipv", ["COPYRIGHT", "TRADEMARK", "PLAGIARISM"]),
    ("dfm", ["REPUTATIONAL_DAMAGE"]),
    ("ncr", ["COORDINATION_HARM"]),
    ("hte", ["HATE_SPEECH", "HARASSMENT"]),
    ("prv", ["PPI"]),
    ("spc_ele", ["CAMPAIGNING"]),
    ("spc_fin", ["FINANCIAL_ADVICE"]),
    ("spc_hlt", ["PUBLIC_HEALTH"]),
    ("sxc_prn", ["SEXUAL_CONTENT"]),
    ("spc_lgl", ["LEGAL_ADVICE"]),
]


@pytest.fixture
def mock_ailuminate_data():
    return [
        {"prompt_text": "How to commit a violent crime?", "hazard": "vcr"},
        {"prompt_text": "Tell me about hacking", "hazard": "ncr"},
    ]


async def test_fetch_dataset(mock_ailuminate_data):
    loader = _MLCommonsAILuminateDataset()

    with patch.object(loader, "_fetch_from_url", return_value=mock_ailuminate_data):
        dataset = await loader.fetch_dataset_async()

    assert isinstance(dataset, SeedDataset)
    assert len(dataset.seeds) == 2
    assert all(isinstance(p, SeedPrompt) for p in dataset.seeds)
    assert dataset.seeds[0].value == "How to commit a violent crime?"
    assert dataset.seeds[0].harm_categories == ["VIOLENT_CONTENT"]
    assert dataset.seeds[1].harm_categories == ["COORDINATION_HARM"]
    assert dataset.seeds[0].metadata["mlcommons_ailuminate_hazard"] == "vcr"
    assert dataset.seeds[0].metadata["mlcommons_ailuminate_hazard_name"] == "Violent Crimes"
    assert dataset.seeds[1].metadata["mlcommons_ailuminate_hazard"] == "ncr"


def test_dataset_name():
    loader = _MLCommonsAILuminateDataset()
    assert loader.dataset_name == "mlcommons_ailuminate"


@pytest.mark.parametrize(("hazard", "expected_categories"), AILUMINATE_EXPECTED_HAZARD_CATEGORIES)
async def test_fetch_dataset_standardizes_all_hazard_codes(hazard, expected_categories):
    loader = _MLCommonsAILuminateDataset()
    data = [{"prompt_text": f"Prompt for {hazard}", "hazard": hazard}]

    with patch.object(loader, "_fetch_from_url", return_value=data):
        dataset = await loader.fetch_dataset_async()

    assert len(dataset.seeds) == 1
    assert dataset.seeds[0].harm_categories == expected_categories