"""Toxicity guardrail for Brazilian Portuguese using ToxiGuardrailPT."""

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

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 = "nicholasKluge/ToxiGuardrailPT"


class ToxicityPTGuardrail(BaseGuardrail):
    """Toxicity detection guardrail for Brazilian Portuguese.

    Uses the ToxiGuardrailPT model (based on BERTimbau) to score text
    for toxicity. The model outputs a single scalar reward score where
    positive values indicate safe content and negative values indicate
    toxic content.

    Both inputs and outputs are checked by default. For output checks,
    the model uses sentence-pair encoding (prompt + response) for more
    accurate contextual scoring.

    Example:
        >>> guard = ToxicityPTGuardrail()
        >>> result = guard.check_input({"message": "Olá, tudo bem?"})
        >>> 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
    # Default overlap between chunks (in tokens)
    _DEFAULT_CHUNK_OVERLAP = 50

    def __init__(
        self,
        name: str = "Toxicity Guardrail (PT-BR)",
        enabled: bool = True,
        threshold: float = 0.0,
        max_length: Optional[int] = 512,
        chunk_overlap: int = 50,
        block_strategy: BlockStrategy = BlockStrategy.RETURN_ERROR_MESSAGE,
        block_message: str = "Requisição bloqueada devido a conteúdo tóxico",
        **config,
    ):
        """Initialize the Toxicity guardrail.

        Args:
            name: Human-readable name for this guardrail.
            enabled: Whether this guardrail is active.
            threshold: Score boundary for toxicity. Texts scoring below this
                value are considered toxic. The model outputs positive scores
                for safe content and negative scores for toxic content, so the
                default of 0.0 is the natural decision boundary.
            max_length: Maximum number of tokens to evaluate. When set to 512
                (default), texts are truncated at the model's limit — fast but
                may miss toxic content in longer texts. 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.max_length = max_length
        self.chunk_overlap = chunk_overlap
        self.block_strategy = block_strategy
        self.block_message = block_message

        logger.debug("Loading ToxiGuardrailPT 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("ToxiGuardrailPT 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",
                },
            )

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

    def check_output(self, output: Any, inputs: Dict[str, Any]) -> GuardrailResult:
        """Check function output for toxic content.

        When inputs are available, uses sentence-pair encoding
        (prompt + response) for contextual scoring.
        """
        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",
                },
            )

        # Build context from inputs for sentence-pair scoring
        input_texts = self._extract_texts(inputs)
        prompt_context = " ".join(input_texts) if input_texts else None

        min_score = self._score_texts(output_texts, prompt_context=prompt_context)
        return self._build_result(
            score=min_score, data_type="output"
        )

    def _score_texts(
        self,
        texts: List[str],
        prompt_context: Optional[str] = None,
    ) -> float:
        """Score a list of texts and return the minimum (most toxic) score.

        Args:
            texts: Texts to score.
            prompt_context: Optional prompt to use as the first segment in
                sentence-pair encoding (for contextual output scoring).

        Returns:
            The minimum score across all texts. Lower = more toxic.
        """
        scores = []
        for text in texts:
            scores.extend(self._score_single_text(text, prompt_context))
        return min(scores)

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

        Returns a list of scores (one per chunk, or just one if truncating).
        """
        if self.max_length is not None:
            # Fast path: truncate to max_length
            return [self._infer(text, prompt_context, truncate=True)]

        # Chunking path: tokenize without truncation, then split into windows
        if prompt_context:
            full = self._tokenizer(
                prompt_context,
                text,
                truncation=False,
                return_token_type_ids=False,
                return_attention_mask=False,
            )
        else:
            full = self._tokenizer(
                text,
                truncation=False,
                return_token_type_ids=False,
                return_attention_mask=False,
            )

        input_ids = full["input_ids"]

        # If it fits in the model window, score directly
        if len(input_ids) <= self.MODEL_MAX_LENGTH:
            return [self._infer(text, prompt_context, truncate=True)]

        # Strip [CLS] and [SEP] added by the tokenizer so we can re-chunk
        # the raw token ids ourselves.
        # BERTimbau adds [CLS] at start and [SEP] at end (and between pairs).
        # We strip them and re-add per chunk.
        raw_ids = input_ids[1:-1]  # remove outer [CLS] ... [SEP]

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

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

            # Rebuild: [CLS] + chunk + [SEP]
            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():
                score = self._model(**tokens).logits.item()
            scores.append(score)

        return scores if scores else [self._infer(text, prompt_context, truncate=True)]

    def _infer(
        self,
        text: str,
        prompt_context: Optional[str] = None,
        truncate: bool = True,
    ) -> float:
        """Run a single forward pass and return the scalar score."""
        if prompt_context:
            tokens = self._tokenizer(
                prompt_context,
                text,
                truncation=truncate,
                max_length=self.MODEL_MAX_LENGTH,
                return_token_type_ids=False,
                return_tensors="pt",
                return_attention_mask=True,
            )
        else:
            tokens = self._tokenizer(
                text,
                truncation=truncate,
                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():
            return self._model(**tokens).logits.item()

    def _build_result(self, score: float, data_type: str) -> GuardrailResult:
        """Build a GuardrailResult from a toxicity score."""
        is_toxic = score < self.threshold

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

        if is_toxic:
            return GuardrailResult(
                action=GuardrailAction.BLOCK,
                metadata={**metadata, "action": "blocked"},
                reason=f"Detected toxic content (score: {score:.3f}, 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,
                "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()]
