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

from unittest.mock import patch

import pytest

from pyrit.datasets.seed_datasets.remote.cbt_bench_dataset import _CBTBenchDataset
from pyrit.models import SeedDataset, SeedPrompt


@pytest.fixture
def mock_cbt_bench_data():
    """Mock data for CBT-Bench dataset."""
    return [
        {
            "id": "1",
            "ori_text": "Original text 1",
            "situation": "I feel guilty for lying to my boyfriend.",
            "thoughts": "I feel ashamed and afraid of losing his trust.",
            "core_belief_fine_grained": ["I am unlovable", "I am immoral"],
        },
        {
            "id": "2",
            "ori_text": "Original text 2",
            "situation": "I am concerned I may have schizophrenia.",
            "thoughts": "I hear voices and see things. I am scared.",
            "core_belief_fine_grained": ["I am powerless, weak, vulnerable", "I am out of control"],
        },
    ]


@pytest.fixture
def mock_cbt_bench_data_missing_thoughts():
    """Mock data with missing thoughts field."""
    return [
        {
            "id": "1",
            "situation": "A situation without thoughts.",
            "thoughts": "",
            "core_belief_fine_grained": ["I am helpless"],
        },
    ]


@pytest.fixture
def mock_cbt_bench_data_empty():
    """Mock data with all empty fields."""
    return [
        {
            "id": "1",
            "situation": "",
            "thoughts": "",
            "core_belief_fine_grained": [],
        },
    ]


class TestCBTBenchDataset:
    """Test the CBT-Bench dataset loader."""

    def test_dataset_name(self):
        """Test dataset_name property."""
        loader = _CBTBenchDataset()
        assert loader.dataset_name == "cbt_bench"

    async def test_fetch_dataset(self, mock_cbt_bench_data):
        """Test fetching CBT-Bench dataset with mocked data."""
        loader = _CBTBenchDataset()

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

            # Check first prompt combines situation and thoughts
            first_prompt = dataset.seeds[0]
            assert "I feel guilty for lying to my boyfriend." in first_prompt.value
            assert "I feel ashamed and afraid of losing his trust." in first_prompt.value
            assert first_prompt.value.startswith("Situation:")
            assert "Thoughts:" in first_prompt.value
            assert first_prompt.data_type == "text"
            assert first_prompt.dataset_name == "cbt_bench"
            assert first_prompt.harm_categories == ["EMOTIONAL", "MENTAL_HEALTH"]
            assert first_prompt.metadata["core_belief_fine_grained"] == ["I am unlovable", "I am immoral"]

    async def test_fetch_dataset_with_custom_config(self, mock_cbt_bench_data):
        """Test fetching with custom HuggingFace config."""
        loader = _CBTBenchDataset(
            source="custom/cbt-bench",
            config="core_major_seed",
        )

        with patch.object(loader, "_fetch_from_huggingface_async", return_value=mock_cbt_bench_data) as mock_fetch:
            dataset = await loader.fetch_dataset_async(cache=False)

            assert len(dataset.seeds) == 2
            mock_fetch.assert_called_once()
            call_kwargs = mock_fetch.call_args.kwargs
            assert call_kwargs["dataset_name"] == "custom/cbt-bench"
            assert call_kwargs["config"] == "core_major_seed"
            assert call_kwargs["split"] == "train"
            assert call_kwargs["cache"] is False

    async def test_fetch_dataset_situation_only(self, mock_cbt_bench_data_missing_thoughts):
        """Test that items with only situation (no thoughts) still work."""
        loader = _CBTBenchDataset()

        with patch.object(loader, "_fetch_from_huggingface_async", return_value=mock_cbt_bench_data_missing_thoughts):
            dataset = await loader.fetch_dataset_async()

            assert len(dataset.seeds) == 1
            assert dataset.seeds[0].value == "A situation without thoughts."

    async def test_fetch_dataset_empty_raises(self, mock_cbt_bench_data_empty):
        """Test that an empty dataset raises ValueError."""
        loader = _CBTBenchDataset()

        with patch.object(loader, "_fetch_from_huggingface_async", return_value=mock_cbt_bench_data_empty):
            with pytest.raises(ValueError, match="SeedDataset cannot be empty"):
                await loader.fetch_dataset_async()

    async def test_fetch_dataset_metadata_includes_config(self, mock_cbt_bench_data):
        """Test that metadata includes the config name."""
        loader = _CBTBenchDataset(config="distortions_seed")

        with patch.object(loader, "_fetch_from_huggingface_async", return_value=mock_cbt_bench_data):
            dataset = await loader.fetch_dataset_async()

            for seed in dataset.seeds:
                assert seed.metadata["config"] == "distortions_seed"

    async def test_fetch_dataset_source_url(self, mock_cbt_bench_data):
        """Test that source URL is correctly set."""
        loader = _CBTBenchDataset()

        with patch.object(loader, "_fetch_from_huggingface_async", return_value=mock_cbt_bench_data):
            dataset = await loader.fetch_dataset_async()

            for seed in dataset.seeds:
                assert seed.source == "https://huggingface.co/datasets/Psychotherapy-LLM/CBT-Bench"