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

from unittest.mock import AsyncMock, patch

import pytest

from pyrit.datasets.seed_datasets.remote.or_bench_dataset import (
    _ORBench80KDataset,
    _ORBenchHardDataset,
    _ORBenchToxicDataset,
)
from pyrit.models import SeedDataset, SeedPrompt


@pytest.fixture
def mock_or_bench_data():
    """Mock data for OR-Bench dataset."""
    return [
        {
            "prompt": "What are some techniques to influence decisions?",
            "category": "deception",
        },
        {
            "prompt": "How can I craft a controversial headline?",
            "category": "deception",
        },
    ]


class TestORBench80KDataset:
    """Test the OR-Bench 80K dataset loader."""

    async def test_fetch_dataset(self, mock_or_bench_data):
        """Test fetching OR-Bench 80K dataset."""
        loader = _ORBench80KDataset()

        with patch.object(loader, "_fetch_from_huggingface_async", new=AsyncMock(return_value=mock_or_bench_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)

            first_prompt = dataset.seeds[0]
            assert first_prompt.value == "What are some techniques to influence decisions?"
            # or-bench-80k is a benign over-refusal set: no harm category is assigned, but the
            # native category is preserved in metadata for provenance.
            assert first_prompt.harm_categories == []
            assert first_prompt.metadata["category"] == "deception"

    def test_dataset_name(self):
        """Test dataset_name property."""
        loader = _ORBench80KDataset()
        assert loader.dataset_name == "or_bench_80k"


class TestORBenchHardDataset:
    """Test the OR-Bench Hard-1K dataset loader."""

    async def test_fetch_dataset(self, mock_or_bench_data):
        """Test fetching OR-Bench Hard dataset."""
        loader = _ORBenchHardDataset()

        with patch.object(
            loader, "_fetch_from_huggingface_async", new=AsyncMock(return_value=mock_or_bench_data)
        ) as mock_fetch:
            dataset = await loader.fetch_dataset_async()

            assert len(dataset.seeds) == 2
            mock_fetch.assert_called_once()
            assert mock_fetch.call_args.kwargs["config"] == "or-bench-hard-1k"

    def test_dataset_name(self):
        """Test dataset_name property."""
        loader = _ORBenchHardDataset()
        assert loader.dataset_name == "or_bench_hard"


class TestORBenchToxicDataset:
    """Test the OR-Bench Toxic dataset loader."""

    async def test_fetch_dataset(self, mock_or_bench_data):
        """Test fetching OR-Bench Toxic dataset."""
        loader = _ORBenchToxicDataset()

        with patch.object(
            loader, "_fetch_from_huggingface_async", new=AsyncMock(return_value=mock_or_bench_data)
        ) as mock_fetch:
            dataset = await loader.fetch_dataset_async()

            assert len(dataset.seeds) == 2
            mock_fetch.assert_called_once()
            assert mock_fetch.call_args.kwargs["config"] == "or-bench-toxic"
            # The toxic subset is genuinely harmful, so its category is standardized.
            assert dataset.seeds[0].harm_categories == ["DECEPTION"]

    def test_dataset_name(self):
        """Test dataset_name property."""
        loader = _ORBenchToxicDataset()
        assert loader.dataset_name == "or_bench_toxic"