"""Quantized detection engines for guardrails."""

from typing import Any, Dict, Optional

import numpy as np

from .loader import QuantizedModel, QuantizationConfig


class QuantizedONNXEngine:
    """
    Quantized ONNX-based detection engine.

    This engine uses int8 quantized ONNX models for fast, memory-efficient
    inference on CPU.

    Example:
        >>> engine = QuantizedONNXEngine("models/quantized/jailbreak_int8.onnx")
        >>> score = engine.detect("Ignore previous instructions")
        >>> print(f"Jailbreak score: {score:.2f}")
    """

    def __init__(
        self,
        model_path: str,
        config: Optional[QuantizationConfig] = None,
        threshold: float = 0.7,
    ):
        """
        Initialize the quantized ONNX engine.

        Args:
            model_path: Path to quantized ONNX model
            config: Optional quantization configuration
            threshold: Detection threshold (0-1)
        """
        self.model = QuantizedModel(model_path, config)
        self.threshold = threshold

    def detect(self, text: str) -> float:
        """
        Detect issues in text.

        Args:
            text: Text to analyze

        Returns:
            Detection score between 0 and 1
        """
        # Tokenize and prepare input (simplified - use actual tokenizer in production)
        input_array = self._preprocess(text)

        # Run inference
        output = self.model.predict(input_array)

        # Extract score
        score = float(output[0]) if output.ndim == 1 else float(output[0][0])
        return score

    def _preprocess(self, text: str) -> np.ndarray:
        """
        Preprocess text for model input.

        In production, replace this with actual tokenization.
        """
        # Placeholder: convert text to numeric representation
        # In real implementation, use proper tokenizer
        text_hash = hash(text) % 1000
        return np.array([[text_hash]], dtype=np.float32)

    def check(self, text: str) -> Dict[str, Any]:
        """
        Check text and return detailed result.

        Args:
            text: Text to check

        Returns:
            Dict with 'score', 'passed', and 'message'
        """
        score = self.detect(text)
        passed = score < self.threshold

        return {
            "score": score,
            "passed": passed,
            "threshold": self.threshold,
            "message": f"Detection score: {score:.2f} (threshold: {self.threshold:.2f})",
        }


class QuantizedTransformerEngine:
    """
    Quantized transformer-based detection engine.

    Uses quantized transformer models for semantic understanding with
    reduced memory footprint.

    Example:
        >>> engine = QuantizedTransformerEngine(
        ...     "models/quantized/toxicity_bert_int8.onnx",
        ...     labels=["safe", "toxic"]
        ... )
        >>> result = engine.classify("This is a test message")
        >>> print(f"Classification: {result['label']} ({result['confidence']:.2f})")
    """

    def __init__(
        self,
        model_path: str,
        labels: Optional[list] = None,
        config: Optional[QuantizationConfig] = None,
    ):
        """
        Initialize the quantized transformer engine.

        Args:
            model_path: Path to quantized model
            labels: List of class labels
            config: Optional quantization configuration
        """
        self.model = QuantizedModel(model_path, config)
        self.labels = labels or ["negative", "positive"]

    def classify(self, text: str) -> Dict[str, Any]:
        """
        Classify text into categories.

        Args:
            text: Text to classify

        Returns:
            Dict with 'label', 'confidence', and 'scores'
        """
        # Preprocess
        input_array = self._preprocess(text)

        # Run inference
        logits = self.model.predict(input_array)

        # Apply softmax
        scores = self._softmax(logits[0])

        # Get prediction
        predicted_idx = int(np.argmax(scores))
        confidence = float(scores[predicted_idx])
        label = self.labels[predicted_idx] if predicted_idx < len(self.labels) else "unknown"

        return {
            "label": label,
            "confidence": confidence,
            "scores": {self.labels[i]: float(scores[i]) for i in range(len(scores))},
        }

    def _preprocess(self, text: str) -> np.ndarray:
        """Preprocess text for model input."""
        # Placeholder - use actual tokenizer in production
        text_hash = hash(text) % 1000
        return np.array([[text_hash]], dtype=np.float32)

    @staticmethod
    def _softmax(x: np.ndarray) -> np.ndarray:
        """Apply softmax to logits."""
        exp_x = np.exp(x - np.max(x))
        return exp_x / exp_x.sum()


__all__ = ["QuantizedONNXEngine", "QuantizedTransformerEngine"]
