"""Prompt injection guardrail using Meta's Prompt Guard 86M model."""

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,
        pipeline,
    )

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


MODEL_ID = "meta-llama/Prompt-Guard-86M"
JAILBREAK_LABEL = "JAILBREAK"


class PromptInjectionGuardrail(BaseGuardrail):
    """Prompt injection detection guardrail using Meta's Prompt Guard 86M model.

    Detects prompt injection attempts in function inputs using a fine-tuned
    sequence classification model. Outputs are not checked since prompt injection
    is an input-side concern.

    Requires a Hugging Face token to download the gated model.
    """

    def __init__(
        self,
        name: str = "Prompt Injection Guardrail",
        enabled: bool = True,
        confidence_threshold: float = 0.5,
        hf_token: Optional[str] = None,
        block_strategy: BlockStrategy = BlockStrategy.RETURN_ERROR_MESSAGE,
        block_message: str = "Request blocked due to prompt injection attempt",
        **config,
    ):
        """Initialize the Prompt Injection guardrail.

        Args:
            name: Human-readable name for this guardrail.
            enabled: Whether this guardrail is active.
            confidence_threshold: Minimum score for the JAILBREAK label to trigger
                a block. Must be between 0.0 and 1.0. Defaults to 0.5.
            hf_token: Hugging Face API token for downloading the gated model.
                Can also be set via the HF_TOKEN environment variable.
            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 Prompt Injection guardrail. "
                "Install with: pip install 'openlayer-guardrails[prompt-injection]'"
            )

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

        self.confidence_threshold = confidence_threshold
        self.hf_token = hf_token
        self.block_strategy = block_strategy
        self.block_message = block_message

        logger.debug("Loading Prompt Guard model: %s", MODEL_ID)
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=hf_token)
        model = AutoModelForSequenceClassification.from_pretrained(
            MODEL_ID, token=hf_token
        )
        self._classifier = pipeline(
            "text-classification",
            model=model,
            tokenizer=tokenizer,
            truncation=True,
            max_length=512,
            device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
        )
        logger.debug("Prompt Guard model loaded successfully")

    def check_input(self, inputs: Dict[str, Any]) -> GuardrailResult:
        """Check function inputs for prompt injection attempts."""
        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",
                },
            )

        # Classify all extracted texts in a single batch call
        raw_results = self._classifier(texts, top_k=1)

        # Find the highest JAILBREAK score across all text chunks
        max_score = 0.0
        for result in raw_results:
            item = result[0]
            if item["label"] == JAILBREAK_LABEL:
                max_score = max(max_score, item["score"])

        is_injection = max_score >= self.confidence_threshold

        metadata = {
            "guardrail": self.name,
            "model_id": MODEL_ID,
            "score": max_score,
            "confidence_threshold": self.confidence_threshold,
            "data_type": "input",
        }

        if is_injection:
            return GuardrailResult(
                action=GuardrailAction.BLOCK,
                metadata={**metadata, "action": "blocked"},
                reason=f"Detected prompt injection attempt (score: {max_score:.3f})",
                block_strategy=self.block_strategy,
                error_message=self.block_message,
            )

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

    def check_output(self, output: Any, inputs: Dict[str, Any]) -> GuardrailResult:
        """Prompt injection is an input-side concern; outputs are always allowed."""
        return GuardrailResult(
            action=GuardrailAction.ALLOW,
            metadata={
                "guardrail": self.name,
                "action": "allow",
                "reason": "output_not_checked",
            },
        )

    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,
                "confidence_threshold": self.confidence_threshold,
                "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()]
