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

import logging
from typing import TYPE_CHECKING, Literal

from typing_extensions import override

from pyrit.datasets.seed_datasets.remote.remote_dataset_loader import (
    _RemoteDatasetLoader,
)
from pyrit.models import Modality, SeedDataset, SeedPrompt

if TYPE_CHECKING:
    from pyrit.models.seeds.seed_group import SeedUnion

logger = logging.getLogger(__name__)


class _XSTestDataset(_RemoteDatasetLoader):
    """
    Loader for the XSTest dataset.

    This dataset contains prompts designed to test exaggerated safety behaviors in language models.

    Reference: [@rottger2023xstest]
    Repository: https://github.com/paul-rottger/exaggerated-safety
    """

    _AUTHORS = [
        "Paul Röttger",
        "Hannah Rose Kirk",
        "Bertie Vidgen",
        "Giuseppe Attanasio",
        "Federico Bianchi",
        "Dirk Hovy",
    ]

    _GROUPS = [
        "Bocconi University",
        "University of Oxford",
        "Stanford University",
    ]

    # Metadata
    modalities: tuple[Modality, ...] = (Modality.TEXT,)
    size: str = "medium"  # 450 safe + unsafe contrast prompts
    tags: frozenset[str] = frozenset({"default", "safety", "refusal"})

    def __init__(
        self,
        *,
        source: str = "https://raw.githubusercontent.com/paul-rottger/exaggerated-safety/a3bb396/xstest_v2_prompts.csv",
        source_type: Literal["public_url", "file"] = "public_url",
    ) -> None:
        """
        Initialize the XSTest dataset loader.

        Args:
            source: URL to the XSTest CSV file. Defaults to the official repository.
            source_type: The type of source ('public_url' or 'file').
        """
        self.source = source
        self.source_type: Literal["public_url", "file"] = source_type

    @property
    @override
    def dataset_name(self) -> str:
        """The dataset name."""
        return "xstest"

    @override
    async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset:
        """
        Fetch XSTest dataset and return as SeedDataset.

        Args:
            cache: Whether to cache the fetched dataset. Defaults to True.

        Returns:
            SeedDataset: A SeedDataset containing the XSTest examples.
        """
        logger.info(f"Loading XSTest dataset from {self.source}")

        examples = self._fetch_from_url(
            source=self.source,
            source_type=self.source_type,
            cache=cache,
        )

        seed_prompts: list[SeedUnion] = [
            SeedPrompt(
                value=example["prompt"],
                data_type="text",
                dataset_name=self.dataset_name,
                # XSTest is an exaggerated-safety (over-refusal) contrast set, so mapping its
                # fields to a harm would mislabel the benign prompts. harm_categories is left
                # empty and the native fields are preserved in metadata.
                harm_categories=[],
                description="A dataset of XSTest examples containing various categories such as violence, drugs, etc.",
                source=self.source,
                metadata={
                    "type": example.get("type", ""),
                    "note": example.get("note", ""),
                    "focus": example.get("focus", ""),
                },
                authors=self._AUTHORS,
                groups=self._GROUPS,
            )
            for example in examples
        ]

        logger.info(f"Successfully loaded {len(seed_prompts)} prompts from XSTest dataset")

        return SeedDataset(seeds=seed_prompts, dataset_name=self.dataset_name)