"""
Starvex ONNX Engine - ML-based toxicity and jailbreak detection

Uses ONNX Runtime with quantized transformer models for fast, accurate
detection on CPU. ~20-50ms latency per check.

Models are cached at ~/.starvex/models/ to avoid re-downloading.

Usage:
    from starvex.guardrails.engines.onnx import OnnxEngine

    engine = OnnxEngine()
    result = engine.check_toxicity("some text", threshold=0.7)
    result = engine.check_jailbreak("some text", threshold=0.7)
"""

import logging
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

# Default model for toxicity detection
DEFAULT_TOXICITY_MODEL = "unitary/unbiased-toxic-roberta"
DEFAULT_JAILBREAK_MODEL = "unitary/unbiased-toxic-roberta"
MODEL_CACHE_DIR = os.path.join(Path.home(), ".starvex", "models")


@dataclass
class OnnxResult:
    """Result from ONNX engine check"""

    passed: bool
    confidence: float = 0.0
    label: str = ""
    message: Optional[str] = None
    details: Dict[str, Any] = field(default_factory=dict)
    scores: Dict[str, float] = field(default_factory=dict)


class OnnxEngine:
    """
    ML-based detection engine using ONNX Runtime.

    Provides high-accuracy toxicity and jailbreak detection using
    quantized transformer models that run on CPU with ~20-50ms latency.

    Features:
    - CPU-only inference via CPUExecutionProvider
    - Quantized models (~50MB) cached at ~/.starvex/models/
    - Thread-safe inference
    - Graceful fallback if ONNX Runtime not installed
    """

    # Toxicity labels from unbiased-toxic-roberta
    TOXICITY_LABELS = [
        "toxicity",
        "severe_toxicity",
        "obscene",
        "threat",
        "insult",
        "identity_attack",
        "sexual_explicit",
    ]

    # Labels that indicate jailbreak-like behavior
    JAILBREAK_INDICATORS = [
        "toxicity",
        "severe_toxicity",
        "threat",
        "identity_attack",
    ]

    def __init__(
        self,
        toxicity_model: str = DEFAULT_TOXICITY_MODEL,
        jailbreak_model: str = DEFAULT_JAILBREAK_MODEL,
        cache_dir: str = MODEL_CACHE_DIR,
        device: str = "cpu",
    ):
        """
        Initialize the ONNX engine.

        Args:
            toxicity_model: HuggingFace model ID for toxicity detection
            jailbreak_model: HuggingFace model ID for jailbreak detection
            cache_dir: Directory to cache downloaded models
            device: Device to run inference on (only "cpu" supported)
        """
        self.toxicity_model_name = toxicity_model
        self.jailbreak_model_name = jailbreak_model
        self.cache_dir = cache_dir
        self.device = device

        self._toxicity_pipeline = None
        self._jailbreak_pipeline = None
        self._available: Optional[bool] = None

    @property
    def is_available(self) -> bool:
        """Check if ONNX Runtime and optimum are available."""
        if self._available is None:
            try:
                import onnxruntime  # noqa: F401
                from optimum.onnxruntime import ORTModelForSequenceClassification  # noqa: F401
                from transformers import AutoTokenizer  # noqa: F401

                self._available = True
            except ImportError:
                self._available = False
                logger.debug(
                    "ONNX Runtime or optimum not installed. Install with: pip install 'starvex[ml]'"
                )
        return self._available

    def _ensure_cache_dir(self) -> None:
        """Create model cache directory if it doesn't exist."""
        os.makedirs(self.cache_dir, exist_ok=True)

    def _load_toxicity_model(self) -> None:
        """Load the toxicity detection model."""
        if self._toxicity_pipeline is not None:
            return

        if not self.is_available:
            raise ImportError("ONNX Runtime not available. Install with: pip install 'starvex[ml]'")

        self._ensure_cache_dir()

        try:
            from optimum.onnxruntime import ORTModelForSequenceClassification
            from transformers import AutoTokenizer, pipeline

            model_path = os.path.join(self.cache_dir, "toxic-roberta-onnx")

            if os.path.exists(model_path):
                logger.debug(f"Loading cached toxicity model from {model_path}")
                model = ORTModelForSequenceClassification.from_pretrained(
                    model_path,
                    provider="CPUExecutionProvider",
                )
                tokenizer = AutoTokenizer.from_pretrained(model_path)
            else:
                logger.info(f"Downloading toxicity model: {self.toxicity_model_name}")
                model = ORTModelForSequenceClassification.from_pretrained(
                    self.toxicity_model_name,
                    export=True,
                    provider="CPUExecutionProvider",
                )
                tokenizer = AutoTokenizer.from_pretrained(self.toxicity_model_name)
                # Cache for next time
                model.save_pretrained(model_path)
                tokenizer.save_pretrained(model_path)
                logger.info(f"Toxicity model cached at {model_path}")

            self._toxicity_pipeline = pipeline(
                "text-classification",
                model=model,
                tokenizer=tokenizer,
                top_k=None,  # Return all labels with scores
                truncation=True,
                max_length=512,
            )
            logger.debug("Toxicity ONNX model loaded successfully")

        except Exception as e:
            logger.error(f"Failed to load toxicity model: {e}")
            self._available = False
            raise

    def _load_jailbreak_model(self) -> None:
        """Load the jailbreak detection model (reuses toxicity model)."""
        # For now, jailbreak detection reuses the toxicity model
        # since toxic-roberta covers threat/attack patterns
        self._load_toxicity_model()
        self._jailbreak_pipeline = self._toxicity_pipeline

    def check_toxicity(
        self,
        text: str,
        threshold: float = 0.7,
    ) -> OnnxResult:
        """
        Check text for toxicity using ML model.

        Args:
            text: Text to analyze
            threshold: Confidence threshold for toxicity (0.0-1.0)

        Returns:
            OnnxResult with toxicity scores and pass/fail
        """
        self._load_toxicity_model()

        try:
            predictions = self._toxicity_pipeline(text)

            # predictions is a list of lists of dicts
            # e.g. [[{"label": "toxicity", "score": 0.95}, ...]]
            if predictions and isinstance(predictions[0], list):
                scores_list = predictions[0]
            else:
                scores_list = predictions

            scores: Dict[str, float] = {}
            for pred in scores_list:
                label = pred.get("label", "").lower()
                score = pred.get("score", 0.0)
                scores[label] = score

            # Get max toxicity score
            toxicity_score = (
                max(scores.get(label, 0.0) for label in self.TOXICITY_LABELS if label in scores)
                if scores
                else 0.0
            )

            is_toxic = toxicity_score >= threshold

            # Find the dominant toxic label
            dominant_label = ""
            if is_toxic:
                toxic_scores = {
                    k: v for k, v in scores.items() if k in self.TOXICITY_LABELS and v >= threshold
                }
                if toxic_scores:
                    dominant_label = max(toxic_scores, key=toxic_scores.get)

            return OnnxResult(
                passed=not is_toxic,
                confidence=toxicity_score,
                label=dominant_label,
                message=(
                    f"Toxic content detected ({dominant_label}: {toxicity_score:.2f})"
                    if is_toxic
                    else "No toxicity detected"
                ),
                details={
                    "threshold": threshold,
                    "backend": "onnx",
                    "model": self.toxicity_model_name,
                },
                scores=scores,
            )

        except Exception as e:
            logger.error(f"ONNX toxicity check error: {e}")
            return OnnxResult(
                passed=True,
                confidence=0.0,
                message=f"ONNX check failed: {e}",
                details={"error": str(e), "backend": "onnx"},
            )

    def check_jailbreak(
        self,
        text: str,
        threshold: float = 0.7,
    ) -> OnnxResult:
        """
        Check text for jailbreak attempts using ML model.

        Uses toxicity model indicators (threat, identity_attack, severe_toxicity)
        as signals for adversarial/jailbreak content.

        Args:
            text: Text to analyze
            threshold: Confidence threshold (0.0-1.0)

        Returns:
            OnnxResult with jailbreak detection result
        """
        self._load_jailbreak_model()

        try:
            predictions = self._jailbreak_pipeline(text)

            if predictions and isinstance(predictions[0], list):
                scores_list = predictions[0]
            else:
                scores_list = predictions

            scores: Dict[str, float] = {}
            for pred in scores_list:
                label = pred.get("label", "").lower()
                score = pred.get("score", 0.0)
                scores[label] = score

            # Check jailbreak indicators
            jailbreak_score = (
                max(
                    scores.get(label, 0.0) for label in self.JAILBREAK_INDICATORS if label in scores
                )
                if scores
                else 0.0
            )

            is_jailbreak = jailbreak_score >= threshold

            return OnnxResult(
                passed=not is_jailbreak,
                confidence=jailbreak_score,
                label="jailbreak" if is_jailbreak else "",
                message=(
                    f"Jailbreak/adversarial content detected (score={jailbreak_score:.2f})"
                    if is_jailbreak
                    else "No jailbreak detected"
                ),
                details={
                    "threshold": threshold,
                    "backend": "onnx",
                    "model": self.jailbreak_model_name,
                    "indicator_scores": {
                        k: v for k, v in scores.items() if k in self.JAILBREAK_INDICATORS
                    },
                },
                scores=scores,
            )

        except Exception as e:
            logger.error(f"ONNX jailbreak check error: {e}")
            return OnnxResult(
                passed=True,
                confidence=0.0,
                message=f"ONNX check failed: {e}",
                details={"error": str(e), "backend": "onnx"},
            )

    @property
    def stats(self) -> Dict[str, Any]:
        """Get engine statistics."""
        return {
            "available": self.is_available,
            "toxicity_model": self.toxicity_model_name,
            "jailbreak_model": self.jailbreak_model_name,
            "cache_dir": self.cache_dir,
            "device": self.device,
            "toxicity_loaded": self._toxicity_pipeline is not None,
            "jailbreak_loaded": self._jailbreak_pipeline is not None,
        }
