"""Gemma-based 5-head multilabel L2 detector.

This detector uses the EmbeddingGemma-300M model with 5 classifier heads:
1. is_threat (binary): benign vs threat
2. threat_family (multiclass, 15): Attack category
3. severity (multiclass, 3): Threat severity level (none, moderate, severe)
4. primary_technique (multiclass, 35): Specific attack technique
5. harm_types (multilabel, 10): Types of potential harm

Model files expected in model directory:
- model_int8.onnx: EmbeddingGemma-300M (INT8 quantized)
- classifier_is_threat_int8.onnx
- classifier_threat_family_int8.onnx
- classifier_severity_int8.onnx
- classifier_primary_technique_int8.onnx
- classifier_harm_types_int8.onnx
- tokenizer.json, config.json, label_config.json, model_metadata.json

NEW: Voting Engine Integration
The ensemble decision logic now uses VotingEngine for transparent weighted voting
instead of boost-based heuristics. Set L2Config.voting.enabled=False to use legacy logic.
"""

from __future__ import annotations

import json
import re
import time
from pathlib import Path
from typing import Any

import numpy as np

from raxe.domain.engine.executor import ScanResult as L1ScanResult
from raxe.domain.ml.gemma_models import (
    GemmaClassificationResult,
    HarmType,
    MultilabelResult,
    PrimaryTechnique,
    Severity,
    ThreatFamily,
)
from raxe.domain.ml.l2_config import L2Config, get_l2_config
from raxe.domain.ml.protocol import L2Prediction, L2Result, L2ThreatType
from raxe.domain.ml.voting import (
    BinaryFirstEngine,
    Decision,
    HeadOutputs,
    VotingResult,
)
from raxe.utils.logging import get_logger

logger = get_logger(__name__)


class GemmaL2Detector:
    """Gemma-based 5-head L2 detector.

    Implements L2Detector protocol with the new Gemma multilabel architecture.

    Performance targets:
    - P95 latency: <50ms (with cache miss)
    - P50 latency: <10ms (with cache hit)
    - Memory: <400MB

    Example:
        detector = GemmaL2Detector(model_dir="/path/to/models")
        result = detector.analyze(text, l1_results)
    """

    DEFAULT_VERSION = "unknown"  # Fallback if model_metadata.json missing
    EMBEDDING_DIM = 768  # Default for model v3, can be overridden by model_metadata.json

    # Handcrafted feature constants for model v3+
    HANDCRAFTED_FEATURE_COUNT = 4  # is_hh_rlhf, text_length, special_char_ratio, question_count

    def __init__(
        self,
        model_dir: str | Path,
        *,
        confidence_threshold: float | None = None,
        harm_thresholds: dict[str, float] | None = None,
        cache_size: int = 1000,
        scorer: Any | None = None,
        l2_config: L2Config | None = None,
    ):
        """Initialize Gemma L2 detector.

        Args:
            model_dir: Path to directory containing ONNX models
            confidence_threshold: Override threshold for binary threat classification
                                  (uses L2Config if not provided)
            harm_thresholds: Per-class thresholds for harm_types multilabel
            cache_size: Size of embedding cache (0 to disable)
            scorer: Optional HierarchicalThreatScorer instance
            l2_config: L2 configuration (uses global config if not provided)
        """
        self.model_dir = Path(model_dir)
        self._l2_config = l2_config or get_l2_config()
        # Use provided threshold, or config threshold
        self.confidence_threshold = (
            confidence_threshold
            if confidence_threshold is not None
            else self._l2_config.thresholds.threat_threshold
        )
        # Use provided harm thresholds, or config thresholds
        self.harm_thresholds = (
            harm_thresholds
            if harm_thresholds is not None
            else self._l2_config.thresholds.harm_type_thresholds.copy()
        )
        self.scorer = scorer

        # Initialize voting engine if enabled
        # Default: BinaryFirstEngine (TPR 90.4%, FPR 7.4%)
        self._voting_enabled = self._l2_config.voting.enabled
        if self._voting_enabled:
            # Use BinaryFirstEngine as default for better TPR/FPR balance
            self._voting_engine = BinaryFirstEngine()
            logger.info(
                "BinaryFirstEngine initialized",
                suppression_quorum=self._voting_engine.config.suppression_quorum,
                enabled=True,
            )
        else:
            self._voting_engine = None
            logger.info("VotingEngine disabled, using legacy ensemble logic")

        # Lazy imports for ONNX runtime
        import onnxruntime as ort
        from transformers import PreTrainedTokenizerFast

        self._ort = ort

        # Load tokenizer from tokenizer.json (portable format)
        logger.info("Loading Gemma tokenizer", model_dir=str(self.model_dir))
        tokenizer_path = self.model_dir / "tokenizer.json"
        if not tokenizer_path.exists():
            raise FileNotFoundError(f"tokenizer.json not found in {self.model_dir}")
        self._tokenizer = PreTrainedTokenizerFast(tokenizer_file=str(tokenizer_path))

        # Configure special tokens from config.json
        config_path = self.model_dir / "config.json"
        if config_path.exists():
            with open(config_path) as f:
                config = json.load(f)
            # Set pad_token_id (Gemma uses id 0 for padding)
            if "pad_token_id" in config:
                self._tokenizer.pad_token_id = config["pad_token_id"]
                # Set pad_token if not already set
                if self._tokenizer.pad_token is None:
                    self._tokenizer.pad_token = self._tokenizer.decode([config["pad_token_id"]])
            if "eos_token_id" in config:
                self._tokenizer.eos_token_id = config["eos_token_id"]
            if "bos_token_id" in config:
                self._tokenizer.bos_token_id = config["bos_token_id"]

        # Load model metadata (single load, reused below)
        self._model_version = self.DEFAULT_VERSION
        self._embedding_dim = self.EMBEDDING_DIM
        self._model_metadata: dict[str, Any] = {}
        metadata_path = self.model_dir / "model_metadata.json"
        if metadata_path.exists():
            try:
                with open(metadata_path) as f:
                    self._model_metadata = json.load(f)
                model_id = self._model_metadata.get("model_id", "unknown")
                model_ver = self._model_metadata.get("model_version", "0.0.0")
                self._model_version = f"{model_id}-v{model_ver}"
                # Read embedding_dim from metadata (single source of truth)
                self._embedding_dim = self._model_metadata.get("embedding_dim", self.EMBEDDING_DIM)
                logger.info(
                    "Loaded model metadata",
                    model_version=self._model_version,
                    embedding_dim=self._embedding_dim,
                )
            except Exception as e:
                logger.warning("Failed to load model_metadata.json", error=str(e))

        # Create session options
        sess_options = ort.SessionOptions()
        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        sess_options.log_severity_level = 3  # ERROR only
        sess_options.intra_op_num_threads = 4
        sess_options.inter_op_num_threads = 1
        sess_options.enable_mem_pattern = True
        sess_options.enable_cpu_mem_arena = True

        providers = ["CPUExecutionProvider"]

        # Load embedding model
        embedding_path = self._find_model_file("model", ".onnx")
        logger.info("Loading embedding model", path=str(embedding_path))
        self._embedding_session = ort.InferenceSession(
            str(embedding_path), sess_options, providers=providers
        )

        # Load classifier heads
        self._classifiers: dict[str, ort.InferenceSession] = {}
        for head in ["is_threat", "threat_family", "severity", "primary_technique", "harm_types"]:
            classifier_path = self._find_model_file(f"classifier_{head}", ".onnx")
            logger.info("Loading classifier", head=head, path=str(classifier_path))
            self._classifiers[head] = ort.InferenceSession(
                str(classifier_path), sess_options, providers=providers
            )

        # Load label config
        label_config_path = self.model_dir / "label_config.json"
        if label_config_path.exists():
            with open(label_config_path) as f:
                self._label_config = json.load(f)
        else:
            self._label_config = {}

        # Load feature scaler for handcrafted features (model v3+)
        # Note: pickle is used here as it's required by sklearn's StandardScaler
        # The scaler file is distributed as part of the trusted model package
        self._feature_scaler = None
        scaler_path = self.model_dir / "feature_scaler.pkl"
        if scaler_path.exists():
            try:
                import pickle

                with open(scaler_path, "rb") as f:
                    self._feature_scaler = pickle.load(f)  # noqa: S301
                logger.info("Loaded feature scaler for handcrafted features")
            except Exception as e:
                logger.warning("Failed to load feature_scaler.pkl", error=str(e))

        # Setup embedding cache
        self._cache_enabled = cache_size > 0
        if self._cache_enabled:
            from raxe.domain.ml.embedding_cache import EmbeddingCache

            self._embedding_cache = EmbeddingCache(max_size=cache_size)
        else:
            self._embedding_cache = None

        logger.info(
            "GemmaL2Detector initialized",
            model_dir=str(self.model_dir),
            embedding_dim=self._embedding_dim,
            cache_size=cache_size,
            confidence_threshold=confidence_threshold,
        )

    def _find_model_file(self, prefix: str, suffix: str) -> Path:
        """Find model file with INT8 preference."""
        # Prefer INT8 quantized version
        int8_path = self.model_dir / f"{prefix}_int8{suffix}"
        if int8_path.exists():
            return int8_path

        # Fall back to non-quantized
        regular_path = self.model_dir / f"{prefix}{suffix}"
        if regular_path.exists():
            return regular_path

        # Try glob pattern
        matches = list(self.model_dir.glob(f"{prefix}*{suffix}"))
        if matches:
            # Prefer INT8 if available
            for match in matches:
                if "int8" in match.name:
                    return match
            return matches[0]

        raise FileNotFoundError(f"No model file found for {prefix}*{suffix} in {self.model_dir}")

    def analyze(
        self,
        text: str,
        l1_results: L1ScanResult,
        context: dict[str, Any] | None = None,
    ) -> L2Result:
        """Analyze text for threats using Gemma 5-head classifier.

        Args:
            text: Text to analyze
            l1_results: Results from L1 rule-based detection
            context: Optional context metadata

        Returns:
            L2Result with predictions from all 5 heads and voting metadata
        """
        start_time = time.perf_counter()

        try:
            # Generate embeddings (with caching) and get token info
            embeddings, token_count, tokens_truncated = self._generate_embeddings(text)

            # Run classification (returns both classification and voting result)
            classification, voting_result = self._classify(embeddings, text=text)

            # Build predictions
            predictions = self._build_predictions(classification, text, voting_result)

            # Apply scorer if available
            hierarchical_score = None
            classification_label = None
            recommended_action = None
            decision_rationale = None
            signal_quality = None

            if self.scorer and classification.is_threat:
                scoring_result = self._apply_scorer(classification, text)
                if scoring_result:
                    hierarchical_score = scoring_result.hierarchical_score
                    classification_label = scoring_result.classification.value
                    recommended_action = scoring_result.action.value
                    decision_rationale = scoring_result.reason
                    signal_quality = {
                        "is_consistent": scoring_result.is_consistent,
                        "variance": scoring_result.variance,
                        "weak_margins_count": scoring_result.weak_margins_count,
                    }

            # Build classification and action from voting result if available
            if voting_result:
                classification_label = self._voting_decision_to_classification(
                    voting_result.decision, voting_result.confidence
                )
                recommended_action = self._voting_decision_to_action(voting_result.decision)
                decision_rationale = f"Voting rule: {voting_result.decision_rule_triggered}"

            duration_ms = (time.perf_counter() - start_time) * 1000

            # Build voting metadata for telemetry
            voting_metadata = voting_result.to_dict() if voting_result else None

            return L2Result(
                predictions=predictions,
                confidence=classification.threat_probability,
                processing_time_ms=duration_ms,
                model_version=self._model_version,
                features_extracted={
                    "text_length": len(text),
                    "l1_detection_count": l1_results.detection_count,
                    "embedding_dim": self._embedding_dim,
                },
                metadata={
                    "detector_type": "gemma",
                    "classification_result": classification.to_dict(),
                    "voting_enabled": self._voting_enabled,
                    # Token count info for telemetry (v2.4)
                    "token_count": token_count,
                    "tokens_truncated": tokens_truncated,
                },
                hierarchical_score=hierarchical_score,
                classification=classification_label,
                recommended_action=recommended_action,
                decision_rationale=decision_rationale,
                signal_quality=signal_quality,
                voting=voting_metadata,
            )

        except Exception as e:
            logger.error("Gemma detection failed", error=str(e), exc_info=True)
            duration_ms = (time.perf_counter() - start_time) * 1000
            return L2Result(
                predictions=[],
                confidence=0.0,
                processing_time_ms=duration_ms,
                model_version=self._model_version,
                metadata={"error": str(e), "detector_type": "gemma"},
            )

    def _voting_decision_to_classification(self, decision: Decision, confidence: float) -> str:
        """Map voting decision to classification label."""
        if decision == Decision.THREAT:
            if confidence >= 0.9:
                return "HIGH_THREAT"
            elif confidence >= 0.75:
                return "THREAT"
            else:
                return "LIKELY_THREAT"
        elif decision == Decision.REVIEW:
            return "REVIEW"
        else:
            return "SAFE"

    def _voting_decision_to_action(self, decision: Decision) -> str:
        """Map voting decision to recommended action."""
        if decision == Decision.THREAT:
            return "BLOCK"
        elif decision == Decision.REVIEW:
            return "MANUAL_REVIEW"
        else:
            return "ALLOW"

    def _generate_embeddings(self, text: str) -> tuple[np.ndarray, int, bool]:
        """Generate embeddings with optional caching.

        The EmbeddingGemma model outputs two tensors:
        - outputs[0]: token embeddings (batch, seq_len, hidden_dim)
        - outputs[1]: pooled embedding (batch, hidden_dim) - used directly

        Returns:
            Tuple of:
            - embeddings: numpy array of shape (1, embedding_dim)
            - token_count: number of tokens after tokenization (max 512)
            - tokens_truncated: True if input was truncated to 512 tokens
        """
        # Always tokenize to get token count (needed for telemetry)
        # This is fast (<1ms) even on cache hit
        inputs = self._tokenizer(
            text,
            padding=True,
            truncation=True,
            max_length=512,
            return_tensors="np",
        )

        # Extract token count info (excluding padding tokens)
        attention_mask = inputs["attention_mask"][0]
        token_count = int(np.sum(attention_mask))

        # Check if truncation occurred by tokenizing without truncation
        # to get the original length
        inputs_no_truncate = self._tokenizer(
            text,
            padding=False,
            truncation=False,
            return_tensors="np",
        )
        original_token_count = len(inputs_no_truncate["input_ids"][0])
        tokens_truncated = original_token_count > 512

        # Check cache
        if self._cache_enabled and self._embedding_cache:
            cached = self._embedding_cache.get(text)
            if cached is not None:
                return cached, token_count, tokens_truncated

        # Get embeddings from model
        outputs = self._embedding_session.run(
            None,
            {
                "input_ids": inputs["input_ids"].astype(np.int64),
                "attention_mask": inputs["attention_mask"].astype(np.int64),
            },
        )

        # outputs[1] is the model's pooled embedding (batch_size, hidden_dim)
        # The model already does internal pooling, so use it directly
        embeddings = outputs[1]

        # Truncate to model's embedding dimension
        embeddings = embeddings[:, : self._embedding_dim]

        # L2 normalize
        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
        embeddings = embeddings / (norms + 1e-9)

        # Cache result
        if self._cache_enabled and self._embedding_cache:
            self._embedding_cache.put(text, embeddings)

        return embeddings, token_count, tokens_truncated

    def _extract_handcrafted_features(self, text: str) -> np.ndarray:
        """Extract handcrafted features for model v3+.

        Features (HANDCRAFTED_FEATURE_COUNT total):
        - is_hh_rlhf: Boolean detecting RLHF-style "Human:" format (word boundary)
        - text_length: Log-normalized text length (log1p(len) / 10.0)
        - special_char_ratio: Ratio of non-word, non-space characters
        - question_count: Count of question marks (divided by 10.0)

        CRITICAL: Feature extraction MUST match training code exactly!
        Training code: raxe-ml/generate_model/scripts/02_train_model.py

        Args:
            text: Input text to extract features from

        Returns:
            Numpy array of shape (1, HANDCRAFTED_FEATURE_COUNT) with scaled features
        """
        # Feature 1: is_hh_rlhf - MUST use \bHuman: pattern to match training
        is_hh_rlhf = 1.0 if re.search(r"\bHuman:", text, re.IGNORECASE) else 0.0

        # Feature 2: text_length - MUST use log1p normalization to match training
        # Training: features[:, 1] = (np.log1p(text_lens) / 10.0)
        text_length = float(np.log1p(len(text)) / 10.0)

        # Feature 3: special_char_ratio - MUST use [^\w\s] regex to match training
        # Training: special_counts = df['text'].str.count(r'[^\w\s]')
        special_char_count = len(re.findall(r"[^\w\s]", text))
        special_char_ratio = special_char_count / max(len(text), 1)

        # Feature 4: question_count - normalized by 10.0 to match training
        question_count = text.count("?") / 10.0

        # Combine features (order must match training)
        features = np.array([[is_hh_rlhf, text_length, special_char_ratio, question_count]])

        # Apply scaler if available
        if self._feature_scaler is not None:
            features = self._feature_scaler.transform(features)

        return features.astype(np.float32)

    def _classify(
        self, embeddings: np.ndarray, text: str | None = None
    ) -> tuple[GemmaClassificationResult, VotingResult | None]:
        """Run all 5 classifier heads with ensemble logic.

        Each classifier returns 2 outputs:
        - [0]: predicted class (int64)
        - [1]: probabilities (float32)

        Args:
            embeddings: Embedding array of shape (1, embedding_dim)
            text: Optional text for handcrafted feature extraction (model v3+)

        Returns:
            Tuple of (GemmaClassificationResult, VotingResult or None)
            VotingResult is None if voting engine is disabled.
        """
        embeddings_f32 = embeddings.astype(np.float32)

        # For model v3+, concatenate embeddings with handcrafted features
        if text is not None and self._feature_scaler is not None:
            handcrafted = self._extract_handcrafted_features(text)
            embeddings_f32 = np.concatenate([embeddings_f32, handcrafted], axis=1)

        # ════════════════════════════════════════════════════════════════════
        # Run all 5 classifier heads (always run all for voting engine)
        # ════════════════════════════════════════════════════════════════════

        # 1. Binary threat classification
        is_threat_outputs = self._classifiers["is_threat"].run(None, {"embeddings": embeddings_f32})
        is_threat_proba = is_threat_outputs[1][0]  # [benign_prob, threat_prob]
        safe_prob = float(is_threat_proba[0])
        threat_prob = float(is_threat_proba[1])

        # 2. Threat family
        family_outputs = self._classifiers["threat_family"].run(
            None, {"embeddings": embeddings_f32}
        )
        family_proba = family_outputs[1][0]
        family_idx = int(np.argmax(family_proba))
        family_confidence = float(family_proba[family_idx])
        threat_family = ThreatFamily.from_index(family_idx)

        # 3. Severity
        severity_outputs = self._classifiers["severity"].run(None, {"embeddings": embeddings_f32})
        severity_proba = severity_outputs[1][0]
        severity_idx = int(np.argmax(severity_proba))
        severity_confidence = float(severity_proba[severity_idx])
        severity = Severity.from_index(severity_idx)

        # 4. Primary technique (always run for voting engine)
        technique_outputs = self._classifiers["primary_technique"].run(
            None, {"embeddings": embeddings_f32}
        )
        technique_proba_arr = technique_outputs[1][0]
        technique_idx = int(np.argmax(technique_proba_arr))
        technique_confidence = float(technique_proba_arr[technique_idx])
        primary_technique = PrimaryTechnique.from_index(technique_idx)
        technique_proba = tuple(float(p) for p in technique_proba_arr)

        # 5. Harm types (always run for voting engine)
        harm_outputs = self._classifiers["harm_types"].run(None, {"embeddings": embeddings_f32})
        harm_proba = harm_outputs[1][0]
        harm_types_result = self._process_multilabel_harm_types(harm_proba)
        harm_max_prob = max(float(p) for p in harm_proba)
        harm_active_labels = [h.value for h in harm_types_result.active_labels]

        # ════════════════════════════════════════════════════════════════════
        # DECISION LOGIC: Voting Engine or Legacy Ensemble
        # ════════════════════════════════════════════════════════════════════

        voting_result: VotingResult | None = None

        if self._voting_enabled and self._voting_engine:
            # Use new VotingEngine for transparent weighted voting
            head_outputs = HeadOutputs(
                binary_threat_prob=threat_prob,
                binary_safe_prob=safe_prob,
                family_prediction=threat_family.value,
                family_confidence=family_confidence,
                severity_prediction=severity.value,
                severity_confidence=severity_confidence,
                technique_prediction=primary_technique.value if primary_technique else None,
                technique_confidence=technique_confidence,
                harm_max_probability=harm_max_prob,
                harm_active_labels=harm_active_labels,
            )

            voting_result = self._voting_engine.vote(head_outputs)

            # Map voting decision to is_threat and effective probability
            is_threat = voting_result.decision in (Decision.THREAT, Decision.REVIEW)
            final_threat_prob = voting_result.confidence
            family_override_triggered = False  # Not used with voting engine
        else:
            # Legacy boost-based ensemble logic
            is_threat, final_threat_prob, family_override_triggered = self._legacy_ensemble_logic(
                threat_prob=threat_prob,
                threat_family=threat_family,
                family_confidence=family_confidence,
                severity=severity,
                primary_technique=primary_technique,
                technique_confidence=technique_confidence,
            )

        # Determine harm_types for result (only include if threat)
        harm_types = harm_types_result if is_threat else None

        result = GemmaClassificationResult(
            is_threat=is_threat,
            threat_probability=final_threat_prob,
            safe_probability=1.0 - final_threat_prob,
            threat_family=threat_family,
            family_confidence=family_confidence,
            family_probabilities=tuple(float(p) for p in family_proba),
            severity=severity,
            severity_confidence=severity_confidence,
            severity_probabilities=tuple(float(p) for p in severity_proba),
            primary_technique=primary_technique,
            technique_confidence=technique_confidence,
            technique_probabilities=technique_proba,
            harm_types=harm_types,
            raw_threat_probability=threat_prob,
            family_override_triggered=family_override_triggered,
        )

        return result, voting_result

    def _legacy_ensemble_logic(
        self,
        threat_prob: float,
        threat_family: ThreatFamily,
        family_confidence: float,
        severity: Severity,
        primary_technique: PrimaryTechnique | None,
        technique_confidence: float,
    ) -> tuple[bool, float, bool]:
        """Legacy boost-based ensemble logic (used when voting is disabled).

        Returns:
            Tuple of (is_threat, effective_threat_prob, family_override_triggered)
        """
        cfg = self._l2_config
        ensemble = cfg.ensemble

        # Start with binary head decision
        effective_threat_prob = threat_prob
        is_threat = threat_prob >= self.confidence_threshold

        # Check for family override
        family_override_triggered = False
        if ensemble.use_family_override:
            family_is_benign = threat_family == ThreatFamily.BENIGN
            family_override_threshold = cfg.thresholds.family_override_threshold

            if not family_is_benign and family_confidence >= family_override_threshold:
                family_override_triggered = True
                is_threat = True
                if effective_threat_prob < self.confidence_threshold:
                    effective_threat_prob = max(
                        effective_threat_prob, self.confidence_threshold + 0.05
                    )

        # Check for always-threat families
        if threat_family.value in ensemble.always_threat_families:
            is_threat = True
            effective_threat_prob = max(effective_threat_prob, 0.75)

        # Apply severity boost
        if ensemble.use_severity_boost:
            if severity == Severity.SEVERE:
                effective_threat_prob = min(
                    1.0, effective_threat_prob + ensemble.severity_boost_amount
                )
                if not is_threat:
                    is_threat = True

        # Apply technique boost
        if ensemble.use_technique_boost and primary_technique:
            if (
                technique_confidence >= ensemble.technique_boost_threshold
                and primary_technique.value in ensemble.high_confidence_techniques
            ):
                effective_threat_prob = min(
                    1.0, effective_threat_prob + ensemble.technique_boost_amount
                )
                if not is_threat:
                    is_threat = True

        return is_threat, effective_threat_prob, family_override_triggered

    def _process_multilabel_harm_types(self, probabilities: np.ndarray) -> MultilabelResult:
        """Process multilabel harm types with per-class thresholds."""
        harm_classes = HarmType.all_classes()
        active_labels = []
        proba_dict: dict[str, float] = {}
        thresholds_dict: dict[str, float] = {}

        for i, harm_type in enumerate(harm_classes):
            prob = float(probabilities[i])
            threshold = self.harm_thresholds.get(harm_type.value, 0.5)
            proba_dict[harm_type.value] = prob
            thresholds_dict[harm_type.value] = threshold

            if prob >= threshold:
                active_labels.append(harm_type)

        return MultilabelResult(
            active_labels=tuple(active_labels),
            probabilities=proba_dict,
            thresholds_used=thresholds_dict,
        )

    def _build_predictions(
        self,
        classification: GemmaClassificationResult,
        text: str,
        voting_result: VotingResult | None = None,
    ) -> list[L2Prediction]:
        """Build L2Predictions from classification result."""
        if not classification.is_threat:
            return []

        # Map family to L2ThreatType
        threat_type = L2ThreatType.from_family(classification.threat_family.value)

        # Build explanation
        explanation_parts = [
            f"Detected {classification.threat_family.value} threat",
            f"with {classification.severity.value} severity",
        ]
        if classification.primary_technique:
            explanation_parts.append(f"using {classification.primary_technique.value} technique")
        if classification.harm_types and classification.harm_types.has_active_labels:
            harm_labels = [h.value for h in classification.harm_types.active_labels]
            explanation_parts.append(f"causing potential harm: {', '.join(harm_labels)}")

        explanation = " ".join(explanation_parts)

        # Build metadata
        # Calculate classification label based on confidence
        if classification.threat_probability >= 0.9:
            classification_label = "HIGH_THREAT"
            action = "BLOCK_ALERT"
        elif classification.threat_probability >= 0.75:
            classification_label = "THREAT"
            action = "BLOCK"
        elif classification.threat_probability >= 0.6:
            classification_label = "LIKELY_THREAT"
            action = "BLOCK_WITH_REVIEW"
        elif classification.threat_probability >= 0.4:
            classification_label = "REVIEW"
            action = "MANUAL_REVIEW"
        else:
            classification_label = "FP_LIKELY"
            action = "ALLOW_WITH_LOG"

        metadata: dict[str, Any] = {
            "is_attack": True,
            "family": classification.threat_family.value,
            "severity": classification.severity.value,
            "classification": classification_label,
            "action": action,
            "risk_score": classification.threat_probability * 100,
            "hierarchical_score": classification.threat_probability,
            "scores": {
                "attack_probability": classification.threat_probability,
                "family_confidence": classification.family_confidence,
                "severity_confidence": classification.severity_confidence,
                "subfamily_confidence": classification.technique_confidence,
            },
        }

        if classification.primary_technique:
            metadata["primary_technique"] = classification.primary_technique.value
            metadata["technique_confidence"] = classification.technique_confidence
            # Alias for CLI formatter compatibility
            metadata["sub_family"] = classification.primary_technique.value

        if classification.harm_types:
            metadata["harm_types"] = classification.harm_types.to_dict()

        # Generate "why_it_hit" explanations
        why_it_hit = self._generate_why_it_hit(classification)
        metadata["why_it_hit"] = why_it_hit

        # Generate recommended actions
        recommended_actions = self._generate_recommended_actions(classification)
        metadata["recommended_action"] = recommended_actions

        # Uncertainty flag
        metadata["uncertain"] = (
            classification.family_confidence < 0.5 or classification.threat_probability < 0.6
        )

        # Family uncertainty flag: binary says threat but family is benign with low confidence
        # This indicates a likely novel attack pattern that doesn't fit known threat families
        family_is_benign = classification.threat_family.value == "benign"
        family_confidence_low = classification.family_confidence < 0.60
        if family_is_benign and family_confidence_low:
            metadata["family_uncertain"] = True
            metadata["threat_type_display"] = "uncategorized_threat"
        else:
            metadata["family_uncertain"] = False

        # Add voting information if available
        if voting_result:
            metadata["voting"] = {
                "decision": voting_result.decision.value,
                "confidence": voting_result.confidence,
                "preset_used": voting_result.preset_used,
                "decision_rule_triggered": voting_result.decision_rule_triggered,
                "threat_vote_count": voting_result.threat_vote_count,
                "safe_vote_count": voting_result.safe_vote_count,
                "abstain_vote_count": voting_result.abstain_vote_count,
                "weighted_ratio": voting_result.weighted_ratio,
            }
            # Update classification and action from voting
            metadata["classification"] = self._voting_decision_to_classification(
                voting_result.decision, voting_result.confidence
            )
            metadata["action"] = self._voting_decision_to_action(voting_result.decision)

        return [
            L2Prediction(
                threat_type=threat_type,
                confidence=classification.threat_probability,
                explanation=explanation,
                features_used=[
                    f"family={classification.threat_family.value}",
                    f"severity={classification.severity.value}",
                    f"technique={classification.primary_technique.value if classification.primary_technique else 'none'}",  # noqa: E501
                    "model=gemma-compact",
                ],
                metadata=metadata,
            )
        ]

    def _generate_why_it_hit(self, classification: GemmaClassificationResult) -> list[str]:
        """Generate human-readable explanations for detection."""
        reasons = []

        # Threat family reason
        family_descriptions = {
            "prompt_injection": "Prompt injection attempting to override instructions",
            "jailbreak": "Jailbreak attempt to bypass safety guidelines",
            "data_exfiltration": "Data exfiltration pattern detected",
            "encoding_or_obfuscation_attack": "Encoded or obfuscated malicious content",
            "rag_or_context_attack": "RAG or context manipulation attack",
            "tool_or_command_abuse": "Tool or command abuse attempt",
            "toxic_or_policy_violating_content": "Toxic or policy-violating content",
            "other_security": "Security-related threat pattern detected",
        }
        family_desc = family_descriptions.get(
            classification.threat_family.value,
            f"{classification.threat_family.value} threat detected",
        )
        reasons.append(family_desc)

        # Severity reason
        if classification.severity == Severity.SEVERE:
            reasons.append(
                f"{classification.severity.value.upper()} severity - immediate action recommended"
            )

        # Technique reason
        has_technique = (
            classification.primary_technique
            and classification.primary_technique != PrimaryTechnique.NONE
        )
        if has_technique:
            technique_desc = classification.primary_technique.value.replace("_", " ")
            reasons.append(f"Specific attack technique: {technique_desc}")

        # Harm types reason
        if classification.harm_types and classification.harm_types.has_active_labels:
            harm_count = classification.harm_types.active_count
            if harm_count == 1:
                harm_label = classification.harm_types.active_labels[0].value.replace("_", " ")
                reasons.append(f"Potential harm type: {harm_label}")
            else:
                reasons.append(f"Multiple harm types detected ({harm_count} categories)")

        return reasons

    def _generate_recommended_actions(self, classification: GemmaClassificationResult) -> list[str]:
        """Generate recommended actions based on classification."""
        actions = []

        if classification.severity == Severity.SEVERE:
            actions.append("SEVERE threat - BLOCK immediately and alert security team")
        elif classification.severity == Severity.MODERATE:
            actions.append("MODERATE severity - Consider blocking or require additional validation")
        else:
            actions.append("Minimal severity - Allow with logging")

        # Family-specific recommendations
        if classification.threat_family == ThreatFamily.DATA_EXFILTRATION:
            actions.append("Review data access controls and implement DLP policies")
        elif classification.threat_family == ThreatFamily.JAILBREAK:
            actions.append("Block and update guardrails to prevent similar attempts")
        elif classification.threat_family == ThreatFamily.PROMPT_INJECTION:
            actions.append("Sanitize input and validate instruction boundaries")
        elif classification.threat_family == ThreatFamily.AGENT_GOAL_HIJACK:
            actions.append("Review agent goals and validate task boundaries")
        elif classification.threat_family == ThreatFamily.PRIVILEGE_ESCALATION:
            actions.append("Audit access controls and review permission boundaries")
        elif classification.threat_family == ThreatFamily.INTER_AGENT_ATTACK:
            actions.append("Isolate agents and validate inter-agent communication")
        elif classification.threat_family == ThreatFamily.MEMORY_POISONING:
            actions.append("Clear memory state and validate memory sources")

        return actions

    def _apply_scorer(self, classification: GemmaClassificationResult, text: str) -> Any:
        """Apply hierarchical threat scorer if available."""
        if not self.scorer:
            return None

        try:
            from raxe.domain.ml.scoring_models import ThreatScore

            # Map to ThreatScore format
            threat_score = ThreatScore(
                binary_threat_score=classification.threat_probability,
                binary_safe_score=classification.safe_probability,
                family_confidence=classification.family_confidence,
                subfamily_confidence=classification.technique_confidence,
                binary_proba=[
                    classification.safe_probability,
                    classification.threat_probability,
                ],
                family_proba=list(classification.family_probabilities),
                subfamily_proba=list(classification.technique_probabilities or [0.0]),
                family_name=classification.threat_family.value,
                subfamily_name=(
                    classification.primary_technique.value
                    if classification.primary_technique
                    else None
                ),
            )

            return self.scorer.score(threat_score, prompt=text)
        except Exception as e:
            logger.warning("Scorer failed", error=str(e))
            return None

    @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()

    @staticmethod
    def _sigmoid(x: np.ndarray) -> np.ndarray:
        """Apply sigmoid to logits."""
        return 1 / (1 + np.exp(-x))

    @property
    def model_info(self) -> dict[str, Any]:
        """Return model information."""
        return {
            "name": "Gemma 5-Head Multilabel Classifier",
            "version": self._model_version,
            "type": "onnx",
            "is_stub": False,
            "size_mb": 300,  # Approximate
            "latency_p95_ms": 50,
            "embedding_model": "google/embeddinggemma-300m",
            "embedding_dim": self._embedding_dim,
            "heads": [
                "is_threat",
                "threat_family",
                "severity",
                "primary_technique",
                "harm_types",
            ],
            "families": [f.value for f in ThreatFamily],
            "description": (
                "Gemma-based 5-head multilabel classifier with EmbeddingGemma-300M "
                "for threat detection, family classification, severity assessment, "
                "technique identification, and harm type prediction."
            ),
        }


def create_gemma_detector(
    model_dir: str | Path,
    *,
    confidence_threshold: float = 0.5,
    harm_thresholds: dict[str, float] | None = None,
    cache_size: int = 1000,
    scorer: Any | None = None,
) -> GemmaL2Detector:
    """Factory function to create Gemma L2 detector.

    Args:
        model_dir: Path to directory containing ONNX models
        confidence_threshold: Threshold for binary threat classification
        harm_thresholds: Per-class thresholds for harm_types multilabel
        cache_size: Size of embedding cache (0 to disable)
        scorer: Optional HierarchicalThreatScorer instance

    Returns:
        GemmaL2Detector instance
    """
    return GemmaL2Detector(
        model_dir=model_dir,
        confidence_threshold=confidence_threshold,
        harm_thresholds=harm_thresholds,
        cache_size=cache_size,
        scorer=scorer,
    )
