"""Toxicity guardrail for English using unitary/toxic-bert."""

import logging
from typing import Any, Dict, List, Optional, Set

from openlayer.lib.guardrails import (
    BaseGuardrail,
    BlockStrategy,
    GuardrailAction,
    GuardrailResult,
)

logger = logging.getLogger(__name__)

try:
    import torch
    from transformers import AutoModelForSequenceClassification, AutoTokenizer

    HAVE_TRANSFORMERS = True
except ImportError:
    HAVE_TRANSFORMERS = False
    torch = None

MODEL_ID = "unitary/toxic-bert"

# Labels in the order the model outputs them (id2label mapping)
LABEL_NAMES = [
    "toxic",
    "severe_toxic",
    "obscene",
    "threat",
    "insult",
    "identity_hate",
]


class ToxicityENGuardrail(BaseGuardrail):
    """Toxicity detection guardrail for English.

    Uses unitary/toxic-bert, a multi-label classifier that scores text
    across six toxicity categories: ``toxic``, ``severe_toxic``,
    ``obscene``, ``threat``, ``insult``, and ``identity_hate``.

    Each category produces an independent probability (0–1) via sigmoid.
    A text is blocked when **any** monitored category exceeds the
    configured threshold.

    Both inputs and outputs are checked.

    Example:
        >>> guard = ToxicityENGuardrail()
        >>> result = guard.check_input({"message": "Have a great day!"})
        >>> result.action
        <GuardrailAction.ALLOW: 'allow'>
    """

    # Maximum sequence length the model supports
    MODEL_MAX_LENGTH = 512
    # Reserve 2 tokens for [CLS] and [SEP]
    _SPECIAL_TOKENS_COUNT = 2

    def __init__(
        self,
        name: str = "Toxicity Guardrail (EN)",
        enabled: bool = True,
        threshold: float = 0.5,
        categories: Optional[Set[str]] = None,
        max_length: Optional[int] = 512,
        chunk_overlap: int = 50,
        block_strategy: BlockStrategy = BlockStrategy.RETURN_ERROR_MESSAGE,
        block_message: str = "Request blocked due to toxic content",
        **config,
    ):
        """Initialize the English Toxicity guardrail.

        Args:
            name: Human-readable name for this guardrail.
            enabled: Whether this guardrail is active.
            threshold: Probability above which a category triggers a block.
                Must be between 0.0 and 1.0. Defaults to 0.5.
            categories: Set of category names to monitor. Defaults to all six
                categories. Valid names: ``toxic``, ``severe_toxic``,
                ``obscene``, ``threat``, ``insult``, ``identity_hate``.
            max_length: Maximum number of tokens to evaluate. When set to 512
                (default), texts are truncated at the model's limit. Set to
                ``None`` to enable chunking mode, which splits long texts
                into overlapping 512-token windows and scores each one.
                Latency scales linearly with the number of chunks.
            chunk_overlap: Number of overlapping tokens between consecutive
                chunks when chunking is enabled. Ignored when ``max_length``
                is set. Defaults to 50.
            block_strategy: How to handle blocked requests.
            block_message: Message returned or raised when a request is blocked.
            **config: Additional configuration passed to BaseGuardrail.
        """
        if not HAVE_TRANSFORMERS:
            raise ImportError(
                "transformers and torch are required for the Toxicity guardrail. "
                "Install with: pip install 'openlayer-guardrails[toxicity]'"
            )

        super().__init__(name=name, enabled=enabled, **config)

        self.threshold = threshold
        self.categories = categories or set(LABEL_NAMES)
        self.max_length = max_length
        self.chunk_overlap = chunk_overlap
        self.block_strategy = block_strategy
        self.block_message = block_message

        # Indices of monitored categories in the model's output
        self._category_indices = [
            i for i, name in enumerate(LABEL_NAMES) if name in self.categories
        ]

        logger.debug("Loading toxic-bert model: %s", MODEL_ID)
        self._tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        self._model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
        self._device = torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )
        self._model.eval()
        self._model.to(self._device)
        logger.debug("toxic-bert model loaded successfully")

    def check_input(self, inputs: Dict[str, Any]) -> GuardrailResult:
        """Check function inputs for toxic content."""
        if not self.enabled:
            return GuardrailResult(
                action=GuardrailAction.ALLOW,
                metadata={
                    "guardrail": self.name,
                    "action": "allow",
                    "reason": "disabled",
                },
            )

        texts = self._extract_texts(inputs)
        if not texts:
            return GuardrailResult(
                action=GuardrailAction.ALLOW,
                metadata={
                    "guardrail": self.name,
                    "action": "allow",
                    "reason": "no_text_content",
                },
            )

        label_scores = self._score_texts(texts)
        return self._build_result(label_scores=label_scores, data_type="input")

    def check_output(self, output: Any, inputs: Dict[str, Any]) -> GuardrailResult:
        """Check function output for toxic content."""
        if not self.enabled:
            return GuardrailResult(
                action=GuardrailAction.ALLOW,
                metadata={
                    "guardrail": self.name,
                    "action": "allow",
                    "reason": "disabled",
                },
            )

        output_texts = self._extract_texts(output)
        if not output_texts:
            return GuardrailResult(
                action=GuardrailAction.ALLOW,
                metadata={
                    "guardrail": self.name,
                    "action": "allow",
                    "reason": "no_text_content",
                },
            )

        label_scores = self._score_texts(output_texts)
        return self._build_result(label_scores=label_scores, data_type="output")

    # ------------------------------------------------------------------
    # Scoring
    # ------------------------------------------------------------------

    def _score_texts(self, texts: List[str]) -> Dict[str, float]:
        """Score texts and return the worst (highest) probability per label.

        When multiple texts or chunks are evaluated, returns the maximum
        probability seen for each category across all of them.
        """
        aggregated: Dict[str, float] = {name: 0.0 for name in LABEL_NAMES}

        for text in texts:
            for chunk_scores in self._score_single_text(text):
                for name, prob in chunk_scores.items():
                    if prob > aggregated[name]:
                        aggregated[name] = prob

        return aggregated

    def _score_single_text(self, text: str) -> List[Dict[str, float]]:
        """Score a single text, chunking if needed.

        Returns a list of dicts (one per chunk) mapping label names to
        probabilities.
        """
        if self.max_length is not None:
            return [self._infer(text)]

        # Chunking path
        full = self._tokenizer(
            text,
            truncation=False,
            return_token_type_ids=False,
            return_attention_mask=False,
        )
        input_ids = full["input_ids"]

        if len(input_ids) <= self.MODEL_MAX_LENGTH:
            return [self._infer(text)]

        # Strip [CLS] … [SEP] so we can re-chunk
        raw_ids = input_ids[1:-1]

        chunk_capacity = self.MODEL_MAX_LENGTH - self._SPECIAL_TOKENS_COUNT
        stride = chunk_capacity - self.chunk_overlap

        results = []
        for start in range(0, len(raw_ids), stride):
            chunk_ids = raw_ids[start : start + chunk_capacity]
            if len(chunk_ids) < 10:
                break

            ids = (
                [self._tokenizer.cls_token_id]
                + chunk_ids
                + [self._tokenizer.sep_token_id]
            )
            attention_mask = [1] * len(ids)

            tokens = {
                "input_ids": torch.tensor([ids], device=self._device),
                "attention_mask": torch.tensor(
                    [attention_mask], device=self._device
                ),
            }

            with torch.no_grad():
                logits = self._model(**tokens).logits
            probs = torch.sigmoid(logits).squeeze(0).tolist()
            results.append(dict(zip(LABEL_NAMES, probs)))

        return results if results else [self._infer(text)]

    def _infer(self, text: str) -> Dict[str, float]:
        """Run a single forward pass and return per-label probabilities."""
        tokens = self._tokenizer(
            text,
            truncation=True,
            max_length=self.MODEL_MAX_LENGTH,
            return_token_type_ids=False,
            return_tensors="pt",
            return_attention_mask=True,
        )
        tokens = {k: v.to(self._device) for k, v in tokens.items()}

        with torch.no_grad():
            logits = self._model(**tokens).logits

        probs = torch.sigmoid(logits).squeeze(0).tolist()
        return dict(zip(LABEL_NAMES, probs))

    # ------------------------------------------------------------------
    # Result building
    # ------------------------------------------------------------------

    def _build_result(
        self, label_scores: Dict[str, float], data_type: str
    ) -> GuardrailResult:
        """Build a GuardrailResult from per-label probabilities."""
        triggered = {
            name: score
            for name, score in label_scores.items()
            if name in self.categories and score >= self.threshold
        }

        metadata = {
            "guardrail": self.name,
            "model_id": MODEL_ID,
            "label_scores": label_scores,
            "threshold": self.threshold,
            "data_type": data_type,
        }

        if triggered:
            worst_label = max(triggered, key=triggered.get)
            return GuardrailResult(
                action=GuardrailAction.BLOCK,
                metadata={
                    **metadata,
                    "action": "blocked",
                    "triggered_categories": triggered,
                },
                reason=(
                    f"Detected toxic content: "
                    f"{', '.join(f'{k}={v:.3f}' for k, v in triggered.items())} "
                    f"(threshold: {self.threshold})"
                ),
                block_strategy=self.block_strategy,
                error_message=self.block_message,
            )

        return GuardrailResult(
            action=GuardrailAction.ALLOW,
            metadata={**metadata, "action": "allow"},
        )

    def get_metadata(self) -> Dict[str, Any]:
        """Get metadata about this guardrail for trace logging."""
        base = super().get_metadata()
        base.update(
            {
                "model_id": MODEL_ID,
                "threshold": self.threshold,
                "categories": sorted(self.categories),
                "max_length": self.max_length,
                "chunking_enabled": self.max_length is None,
                "transformers_available": HAVE_TRANSFORMERS,
            }
        )
        return base

    def _extract_texts(self, data: Any) -> List[str]:
        """Recursively extract text strings from arbitrary data structures."""
        texts = []
        if isinstance(data, str):
            texts.append(data)
        elif isinstance(data, dict):
            for value in data.values():
                texts.extend(self._extract_texts(value))
        elif isinstance(data, (list, tuple)):
            for item in data:
                texts.extend(self._extract_texts(item))
        elif hasattr(data, "__str__") and not isinstance(data, (int, float, bool)):
            text_repr = str(data)
            if text_repr and text_repr not in ("True", "False", "None"):
                texts.append(text_repr)
        return [t for t in texts if t and t.strip()]
