"""Active learning engine for fix quality improvement.

Uses Thompson Sampling for strategy selection and River for
online learning of fix success prediction.
"""

from __future__ import annotations

import json
import logging
import random
from collections import deque
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Protocol, Tuple

if TYPE_CHECKING:
    from autodebug.learning.cross_org_aggregator import CrossOrgAggregator
    from autodebug.learning.pattern_matcher import PatternMatch, PatternMatcher
    from autodebug.strategies.bandit_selector import ThompsonSamplingSelector

logger = logging.getLogger(__name__)


from river import linear_model, preprocessing
from river import metrics as river_metrics


@dataclass
class LearningSignal:
    """A learning signal from user feedback.

    Captures all information needed to update the learning models.
    """

    fix_id: str
    exception_type: str
    exception_message: str
    strategy_used: str
    outcome: str  # "accepted", "rejected", "modified", "applied_worked", "applied_failed"
    user_modifications: Optional[str] = None
    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    features: Dict[str, Any] = field(default_factory=dict)

    @property
    def is_success(self) -> bool:
        """Whether this outcome is considered a success."""
        return self.outcome in ("accepted", "applied_worked", "modified")

    @property
    def is_failure(self) -> bool:
        """Whether this outcome is considered a failure."""
        return self.outcome in ("rejected", "applied_failed")


class DatabaseProtocol(Protocol):
    """Protocol for database interactions."""

    def store_feedback(
        self,
        fix_id: str,
        outcome: str,
        timestamp: datetime,
        modifications: Optional[str] = None,
    ) -> None:
        """Store feedback event."""
        ...

    def get_fix(self, fix_id: str) -> Optional[Dict[str, Any]]:
        """Get fix by ID."""
        ...

    def store_example(self, example: Dict[str, Any], embedding: List[float]) -> None:
        """Store a learned example."""
        ...

    def search_examples(
        self,
        embedding: List[float],
        filter: Dict[str, Any],
        limit: int,
    ) -> List[Dict[str, Any]]:
        """Search for similar examples."""
        ...

    def count_rejections(self, key: str, days: int) -> int:
        """Count rejections for a pattern."""
        ...

    def create_alert(self, alert: Dict[str, Any]) -> None:
        """Create an alert for review."""
        ...


class RiverOnlineLearner:
    """Online learning model for fix quality prediction using River.

    F-08: Supports both linear (LogisticRegression) and tree-based
    (HoeffdingAdaptiveTreeClassifier) models. The adaptive tree can
    capture non-linear feature interactions and has built-in ADWIN
    drift detection, addressing both F-08 and partially F-01.
    """

    # Decay rates: aggressive decay when drift detected, slow decay otherwise
    DRIFT_DECAY_RATE = 0.85  # Aggressive: recent samples matter much more
    NO_DRIFT_DECAY_RATE = 0.995  # Very slow: stable distribution, gentle forgetting

    def __init__(self, use_tree: bool = False):
        """Initialize the online learner.

        Args:
            use_tree: If True, use HoeffdingAdaptiveTreeClassifier
                (captures non-linear interactions, has built-in drift detection).
                If False (default), use LogisticRegression (linear, reliable
                with small sample sizes). Switch to tree when you have >50
                samples and need non-linear pattern capture.
        """
        from river import optim

        self._use_tree = use_tree

        if use_tree:
            try:
                from river import tree as river_tree

                self._model = river_tree.HoeffdingAdaptiveTreeClassifier(
                    grace_period=20,
                    max_depth=6,
                    leaf_prediction="nba",
                )
            except (ImportError, AttributeError):
                # Fallback to linear model if tree not available
                self._use_tree = False
                self._model = (
                    preprocessing.StandardScaler()
                    | linear_model.LogisticRegression(optimizer=optim.SGD(0.01))
                )
        else:
            self._model = (
                preprocessing.StandardScaler()
                | linear_model.LogisticRegression(optimizer=optim.SGD(0.01))
            )

        self._accuracy = river_metrics.Accuracy()
        self._log_loss = river_metrics.LogLoss()
        self._sample_count = 0
        # Manual rolling window for recent accuracy
        self._rolling_window_size = 100
        self._rolling_window: deque = deque(maxlen=self._rolling_window_size)
        # Effective sample weight, decayed over time based on drift state
        self._current_weight = 1.0

    def learn(
        self,
        features: Dict[str, float],
        success: bool,
        drift_detected: bool = False,
    ) -> None:
        """Incrementally update model with new observation.

        Temporal decay is drift-aware: aggressive decay when ADWIN
        signals drift, very slow decay under stable distribution.

        Args:
            features: Feature dict for this example
            success: Whether the fix succeeded
            drift_detected: Whether concept drift was detected on this step
        """
        if not features:
            return

        self._sample_count += 1

        # Apply drift-aware temporal decay to sample weight
        if drift_detected:
            self._current_weight *= self.DRIFT_DECAY_RATE
        else:
            self._current_weight *= self.NO_DRIFT_DECAY_RATE

        # Predict first (for accuracy tracking)
        pred = self._model.predict_one(features)
        if pred is not None:
            correct = (pred > 0.5) == success
            self._accuracy.update(success, pred > 0.5)
            self._rolling_window.append(correct)

        # Track calibration via log loss
        pred_proba = self._model.predict_proba_one(features)
        if pred_proba:
            self._log_loss.update(success, pred_proba)

        # Learn from this example with drift-aware weight applied
        # River's learn_one supports a sample_weight parameter
        self._model.learn_one(features, success, w=self._current_weight)

    def predict_success(self, features: Dict[str, float]) -> float:
        """Predict success probability for a fix.

        Args:
            features: Feature dict for prediction

        Returns:
            Predicted success probability (0-1)
        """
        if not features:
            return 0.5

        pred = self._model.predict_proba_one(features)
        if pred:
            return pred.get(True, 0.5)
        return 0.5

    def get_accuracy(self) -> float:
        """Get current model accuracy.

        Returns:
            Accuracy on seen examples
        """
        return self._accuracy.get()

    def get_rolling_accuracy(self) -> float:
        """Get rolling accuracy over recent predictions."""
        if not self._rolling_window:
            return 0.0
        return sum(self._rolling_window) / len(self._rolling_window)

    def get_log_loss(self) -> float:
        """Get log loss for calibration tracking."""
        return self._log_loss.get()

    @property
    def sample_count(self) -> int:
        """Number of samples learned from."""
        return self._sample_count


class PerformanceMonitor:
    """Monitors learning performance and triggers fallback when quality drops.

    Tracks rolling acceptance rate and falls back to static priority
    ordering when the learned model underperforms. Uses River ML's ADWIN
    detector for statistically rigorous concept drift detection.
    """

    def __init__(
        self,
        window_size: int = 50,
        fallback_threshold: float = 0.3,
        recovery_threshold: float = 0.4,
        drift_detector: str = "adwin",
    ):
        """Initialize the performance monitor.

        Args:
            window_size: Number of recent fixes to track
            fallback_threshold: Rate below which fallback triggers
            recovery_threshold: Rate above which learning resumes
        """
        self.window_size = window_size
        self.fallback_threshold = fallback_threshold
        self.recovery_threshold = recovery_threshold
        self._recent_outcomes: deque = deque(maxlen=window_size)
        self._in_fallback = False
        self._fallback_count = 0

        # ADWIN-based concept drift detection (River ML)
        self._drift_detector = None
        self._drift_detector_name = "none"
        self._drift_detected = False
        self._drift_count = 0
        self._last_drift_at: Optional[datetime] = None
        self._on_drift_callback = None
        try:
            from river import drift as river_drift

            detector_name = (drift_detector or "adwin").strip().lower()
            if detector_name in ("pagehinkley", "page_hinkley"):
                self._drift_detector = river_drift.PageHinkley()
                self._drift_detector_name = "PageHinkley"
            else:
                self._drift_detector = river_drift.ADWIN()
                self._drift_detector_name = "ADWIN"
        except ImportError:
            logger.debug("River ML drift detection unavailable; using threshold-only monitoring")

    def record_outcome(self, success: bool) -> None:
        """Record a fix outcome."""
        self._recent_outcomes.append(success)

        # Feed outcome into ADWIN drift detector
        self._drift_detected = False
        if self._drift_detector is not None:
            update_result = self._drift_detector.update(1.0 if success else 0.0)
            if bool(getattr(self._drift_detector, "drift_detected", False) or update_result is True):
                self._drift_detected = True
                self._drift_count += 1
                self._last_drift_at = datetime.now(timezone.utc)
                logger.warning(
                    "Concept drift detected "
                    f"({self._drift_detector_name}) "
                    f"(detection #{self._drift_count}, "
                    f"current rate: {self.rolling_acceptance_rate:.0%})"
                )
                if self._on_drift_callback is not None:
                    try:
                        event = {
                            "detector": self._drift_detector_name,
                            "drift_count": self._drift_count,
                            "acceptance_rate": self.rolling_acceptance_rate,
                            "detected_at": self._last_drift_at.isoformat() if self._last_drift_at else None,
                        }
                        self._on_drift_callback(event)
                    except TypeError:
                        # Backward-compatible callback signature: callback()
                        self._on_drift_callback()
                    except Exception as e:
                        logger.debug(f"Drift callback error: {e}")

        # Check if we should trigger or recover from fallback
        rate = self.rolling_acceptance_rate
        if not self._in_fallback and rate < self.fallback_threshold and len(self._recent_outcomes) >= 10:
            self._in_fallback = True
            self._fallback_count += 1
            logger.warning(
                f"Learning fallback triggered: acceptance rate {rate:.0%} "
                f"below threshold {self.fallback_threshold:.0%}"
            )
        elif self._in_fallback and rate >= self.recovery_threshold:
            self._in_fallback = False
            logger.info(
                f"Learning recovered: acceptance rate {rate:.0%} "
                f"above recovery threshold {self.recovery_threshold:.0%}"
            )

    @property
    def rolling_acceptance_rate(self) -> float:
        """Get rolling acceptance rate."""
        if not self._recent_outcomes:
            return 0.5
        return sum(self._recent_outcomes) / len(self._recent_outcomes)

    @property
    def in_fallback(self) -> bool:
        """Whether currently in fallback mode."""
        return self._in_fallback

    @property
    def fallback_count(self) -> int:
        """Number of times fallback has been triggered."""
        return self._fallback_count

    @property
    def drift_detected(self) -> bool:
        """Whether concept drift was detected on the most recent outcome."""
        return self._drift_detected

    @property
    def drift_count(self) -> int:
        """Total number of concept drift events detected."""
        return self._drift_count

    @property
    def drift_detector_name(self) -> str:
        """Name of the active drift detector."""
        return self._drift_detector_name

    @property
    def last_drift_at(self) -> Optional[datetime]:
        """Timestamp of the most recent drift event."""
        return self._last_drift_at

    @property
    def on_drift_callback(self):
        """Callback invoked when ADWIN detects concept drift."""
        return self._on_drift_callback

    @on_drift_callback.setter
    def on_drift_callback(self, callback) -> None:
        self._on_drift_callback = callback


class SimplePersistentLearner:
    """Simple learner that persists to a JSON file.

    Used as a lightweight alternative when a full database
    is not available.
    """

    def __init__(self, path: Optional[str] = None):
        """Initialize the learner.

        Args:
            path: Path to persistence file
        """
        self.path = Path(path).expanduser() if path else None
        self.examples: List[Dict[str, Any]] = []
        self.alerts: List[Dict[str, Any]] = []
        self.rejections: Dict[str, List[datetime]] = {}
        self._load()

    def store_example(self, example: Dict[str, Any]) -> None:
        """Store a learned example."""
        self.examples.append(example)
        self._save()

    def get_examples(
        self,
        exception_type: Optional[str] = None,
        limit: int = 10,
    ) -> List[Dict[str, Any]]:
        """Get stored examples.

        Args:
            exception_type: Filter by exception type
            limit: Maximum examples to return

        Returns:
            List of examples
        """
        filtered = self.examples
        if exception_type:
            filtered = [e for e in filtered if e.get("exception_type") == exception_type]
        return filtered[-limit:]

    def record_rejection(self, key: str) -> None:
        """Record a rejection for tracking repeated failures."""
        if key not in self.rejections:
            self.rejections[key] = []
        self.rejections[key].append(datetime.now(timezone.utc))
        self._save()

    def count_recent_rejections(self, key: str, days: int = 7) -> int:
        """Count recent rejections for a key."""
        if key not in self.rejections:
            return 0
        cutoff = datetime.now(timezone.utc).replace(
            hour=0, minute=0, second=0, microsecond=0
        )
        from datetime import timedelta

        cutoff = cutoff - timedelta(days=days)
        return sum(1 for ts in self.rejections[key] if ts >= cutoff)

    def create_alert(self, alert: Dict[str, Any]) -> None:
        """Create an alert for review."""
        self.alerts.append(alert)
        self._save()

    def _load(self) -> None:
        """Load persisted data."""
        if not self.path or not self.path.exists():
            return

        try:
            data = json.loads(self.path.read_text())
            self.examples = data.get("examples", [])
            self.alerts = data.get("alerts", [])

            # Parse rejection timestamps
            raw_rejections = data.get("rejections", {})
            self.rejections = {
                k: [datetime.fromisoformat(ts) for ts in v]
                for k, v in raw_rejections.items()
            }
        except Exception as e:
            logger.debug(f"Failed to load learner data: {e}")

    def _save(self) -> None:
        """Save data to persistence file using atomic write."""
        if not self.path:
            return

        try:
            import os
            import tempfile

            self.path.parent.mkdir(parents=True, exist_ok=True)

            # Serialize rejection timestamps
            serialized_rejections = {
                k: [ts.isoformat() for ts in v]
                for k, v in self.rejections.items()
            }

            data = {
                "examples": self.examples[-1000:],  # Keep last 1000
                "alerts": self.alerts[-100:],  # Keep last 100
                "rejections": serialized_rejections,
            }
            content = json.dumps(data, indent=2, default=str)

            # Atomic write: write to temp file, then rename
            fd, tmp_path = tempfile.mkstemp(
                dir=str(self.path.parent),
                suffix=".tmp",
            )
            try:
                with os.fdopen(fd, 'w') as f:
                    f.write(content)
                os.replace(tmp_path, str(self.path))
            except BaseException:
                # Clean up temp file on failure
                try:
                    os.unlink(tmp_path)
                except OSError:
                    pass
                raise
        except Exception as e:
            logger.debug(f"Failed to save learner data: {e}")


class ActiveLearner:
    """Main active learning coordinator.

    Integrates Thompson Sampling for strategy selection with
    online learning for fix quality prediction. Optionally integrates
    with cross-organization learning for global pattern matching.
    """

    # Reward weights for different outcomes (used with Thompson Sampling)
    OUTCOME_REWARDS = {
        "accepted": 1.0,
        "modified": 0.7,
        "applied_worked": 1.0,
        "applied_failed": 0.0,
        "rejected": 0.0,
    }

    def __init__(
        self,
        thompson_selector: Optional[ThompsonSamplingSelector] = None,
        persistence_path: Optional[str] = None,
        cross_org: Optional[CrossOrgAggregator] = None,
        org_id: Optional[str] = None,
    ):
        """Initialize the active learner.

        Args:
            thompson_selector: Thompson Sampling selector for strategy learning
            persistence_path: Path to persist learning data
            cross_org: Cross-organization aggregator for global patterns
            org_id: Organization identifier for pattern contribution
        """
        self._thompson = thompson_selector
        self.predictor = RiverOnlineLearner()
        self.persistence = SimplePersistentLearner(persistence_path)
        self.monitor = PerformanceMonitor()

        # Wire ADWIN drift detection to trigger model retraining
        self.monitor.on_drift_callback = self._on_concept_drift

        # Cross-org learning integration
        self._cross_org = cross_org
        self._org_id = org_id
        self._pattern_matcher: Optional[PatternMatcher] = None

        # Track feedback counts
        self.feedback_counts: Dict[str, int] = {
            "accepted": 0,
            "rejected": 0,
            "modified": 0,
            "applied_worked": 0,
            "applied_failed": 0,
        }

        # Track recent fix IDs for deduplication
        self._recent_fix_ids: Dict[str, datetime] = {}

        # Replay buffer for catastrophic forgetting protection (reservoir sampling)
        self._replay_buffer: List[Tuple[Dict[str, float], bool]] = []
        self._replay_buffer_capacity = 100
        self._replay_total_seen = 0

    @property
    def thompson(self) -> Optional[ThompsonSamplingSelector]:
        """Get or create Thompson Sampling selector."""
        if self._thompson is None:
            try:
                from autodebug.strategies.bandit_selector import ThompsonSamplingSelector

                self._thompson = ThompsonSamplingSelector()
            except ImportError:
                pass
        return self._thompson

    def _add_to_replay_buffer(
        self, features: Dict[str, float], success: bool
    ) -> None:
        """Add a sample to the replay buffer using reservoir sampling.

        Maintains a fixed-size buffer of past important examples.
        Uses Algorithm R (Vitter) for uniform reservoir sampling so
        every past sample has equal probability of being retained.

        Args:
            features: Feature dict for this example
            success: Whether the fix succeeded
        """
        self._replay_total_seen += 1
        if len(self._replay_buffer) < self._replay_buffer_capacity:
            self._replay_buffer.append((features, success))
        else:
            # Reservoir sampling: replace with probability capacity/total_seen
            idx = random.randint(0, self._replay_total_seen - 1)
            if idx < self._replay_buffer_capacity:
                self._replay_buffer[idx] = (features, success)

    def _replay_buffer_samples(self, drift_detected: bool = False) -> None:
        """Replay a mini-batch from the buffer to prevent catastrophic forgetting.

        Periodically replays old examples so the model retains
        knowledge of previously learned patterns.

        Args:
            drift_detected: Whether concept drift was detected on this step
        """
        if not self._replay_buffer:
            return

        # Replay a small random subset (up to 5 samples) every 10 new samples
        if self._replay_total_seen % 10 != 0:
            return

        batch_size = min(5, len(self._replay_buffer))
        batch = random.sample(self._replay_buffer, batch_size)
        for features, success in batch:
            self.predictor.learn(features, success, drift_detected=drift_detected)

    # F-11: Valid temporal ordering of feedback signals.
    # Defines which outcomes must precede others for the same fix_id.
    _OUTCOME_PREREQUISITES = {
        "applied_worked": {"accepted", "modified"},
        "applied_failed": {"accepted", "modified"},
    }

    # F-11: Contradictory outcome pairs that cannot both be valid
    _CONTRADICTORY_OUTCOMES = {
        frozenset({"accepted", "rejected"}),
        frozenset({"applied_worked", "applied_failed"}),
    }

    def _validate_feedback(self, signal: LearningSignal) -> bool:
        """Validate feedback signal for noise robustness.

        Rejects signals with impossible combinations, deduplicates, and
        checks temporal ordering (F-11: contradictory feedback detection).

        Args:
            signal: Learning signal to validate

        Returns:
            True if signal is valid
        """
        # Reject signals with missing required fields
        if not signal.strategy_used or not signal.exception_type:
            logger.debug("Rejecting signal: missing strategy or exception type")
            return False

        # Deduplicate: ignore same fix_id+outcome within 60 seconds
        now = datetime.now(timezone.utc)
        dedup_key = f"{signal.fix_id}:{signal.outcome}"
        if dedup_key in self._recent_fix_ids:
            last_seen = self._recent_fix_ids[dedup_key]
            from datetime import timedelta
            if (now - last_seen) < timedelta(seconds=60):
                logger.debug(f"Rejecting duplicate signal for {signal.fix_id}")
                return False

        # F-11: Check temporal ordering -- certain outcomes require
        # prerequisites (e.g., "applied_worked" must follow "accepted")
        required_priors = self._OUTCOME_PREREQUISITES.get(signal.outcome)
        if required_priors:
            prior_outcomes = {
                key.split(":", 1)[1]
                for key in self._recent_fix_ids
                if key.startswith(f"{signal.fix_id}:")
            }
            if not prior_outcomes & required_priors:
                logger.debug(
                    f"Rejecting signal for {signal.fix_id}: "
                    f"'{signal.outcome}' requires prior {required_priors}, "
                    f"but only found {prior_outcomes or 'none'}"
                )
                return False

        # F-11: Reject contradictory outcomes for the same fix_id
        prior_outcomes_for_fix = {
            key.split(":", 1)[1]
            for key in self._recent_fix_ids
            if key.startswith(f"{signal.fix_id}:")
        }
        for pair in self._CONTRADICTORY_OUTCOMES:
            if signal.outcome in pair and (pair - {signal.outcome}) & prior_outcomes_for_fix:
                contradicting = (pair - {signal.outcome}) & prior_outcomes_for_fix
                logger.debug(
                    f"Rejecting contradictory signal for {signal.fix_id}: "
                    f"'{signal.outcome}' contradicts prior {contradicting}"
                )
                return False

        self._recent_fix_ids[dedup_key] = now

        # Clean old entries (keep last 1000)
        if len(self._recent_fix_ids) > 1000:
            sorted_keys = sorted(
                self._recent_fix_ids, key=lambda k: self._recent_fix_ids[k]
            )
            for k in sorted_keys[:len(sorted_keys) - 500]:
                del self._recent_fix_ids[k]

        return True

    def _on_concept_drift(self, event: Optional[Dict[str, Any]] = None) -> None:
        """Handle concept drift detected by ADWIN.

        Resets the online predictor so it can re-learn from the new
        data distribution, rather than carrying stale weights.
        """
        logger.warning(
            "Concept drift detected — resetting online predictor "
            f"(detector={event.get('detector') if event else 'unknown'}, "
            f"had {self.predictor.sample_count} samples)"
        )
        self.predictor = RiverOnlineLearner()

    def record_feedback(
        self,
        signal: LearningSignal,
        stack_trace: Optional[List[str]] = None,
    ) -> None:
        """Record feedback and trigger learning updates.

        Args:
            signal: Learning signal with feedback details
            stack_trace: Optional stack trace for cross-org contribution
        """
        # Validate signal for noise robustness
        if not self._validate_feedback(signal):
            return

        # Update feedback counts
        if signal.outcome in self.feedback_counts:
            self.feedback_counts[signal.outcome] += 1

        # Update performance monitor
        self.monitor.record_outcome(signal.is_success)

        # Update Thompson Sampling with reward shaping
        if self.thompson and signal.strategy_used:
            outcome = signal.outcome if signal.outcome in self.OUTCOME_REWARDS else None
            self.thompson.update(
                signal.exception_type,
                signal.strategy_used,
                signal.is_success,
                outcome=outcome,
            )

            # Negative feedback weighting: extra penalty for repeated rejections
            if signal.is_failure:
                key = f"{signal.exception_type}:{signal.strategy_used}"
                recent_rejects = self.persistence.count_recent_rejections(key, days=7)
                if recent_rejects > 3:
                    # Apply extra negative update for persistent failures
                    self.thompson.update(
                        signal.exception_type,
                        signal.strategy_used,
                        False,
                    )

        # Update online predictor (skip if in fallback mode)
        if not self.monitor.in_fallback:
            features = self._extract_features(signal)
            if features:
                # Uncertainty sampling: predict before learning.
                # High-uncertainty samples (near 0.5) get an extra learning
                # pass so the model focuses on its decision boundary.
                pred_prob = self.predictor.predict_success(features)
                uncertainty = self._compute_uncertainty(pred_prob)

                self.predictor.learn(
                    features,
                    signal.is_success,
                    drift_detected=self.monitor.drift_detected,
                )

                # If uncertainty is high (>0.7), learn from this sample
                # a second time to prioritize boundary examples
                if uncertainty > 0.7:
                    self.predictor.learn(
                        features,
                        signal.is_success,
                        drift_detected=self.monitor.drift_detected,
                    )

                # Replay buffer: store sample and periodically replay
                # old examples to prevent catastrophic forgetting
                self._add_to_replay_buffer(features, signal.is_success)
                self._replay_buffer_samples(
                    drift_detected=self.monitor.drift_detected,
                )

        # If accepted, store as positive example
        if signal.outcome == "accepted":
            self._store_positive_example(signal)

            # Contribute to cross-org learning if enabled
            if self._cross_org and self._org_id and stack_trace:
                self.contribute_to_global(
                    exception_type=signal.exception_type,
                    exception_message=signal.exception_message,
                    stack_trace=stack_trace,
                    strategy_used=signal.strategy_used,
                    was_accepted=True,
                )

        # Track rejections for pattern detection
        if signal.outcome == "rejected":
            self._track_rejection(signal)

            # Also contribute rejections for negative learning
            if self._cross_org and self._org_id and stack_trace:
                self.contribute_to_global(
                    exception_type=signal.exception_type,
                    exception_message=signal.exception_message,
                    stack_trace=stack_trace,
                    strategy_used=signal.strategy_used,
                    was_accepted=False,
                )

    # F-09: Known high-value exception types and strategies are kept as
    # explicit features for interpretability.  All other (rare) values are
    # mapped into a fixed set of hash buckets to bound feature-space growth.
    _KNOWN_EXCEPTIONS = frozenset([
        "keyerror", "attributeerror", "typeerror", "valueerror",
        "indexerror", "nameerror", "importerror", "assertionerror",
        "runtimeerror", "oserror", "ioerror", "filenotfounderror",
        "connectionerror", "timeouterror", "jsondecodeerror",
    ])
    _KNOWN_STRATEGIES = frozenset([
        "use_get", "use_get_default", "add_null_check", "fix_typo",
        "add_error_handling", "fix_the_code", "fix_the_test",
        "add_type_conversion", "fix_import_path", "add_retry",
        "fix_argument_order", "fix_value_source", "add_default_value",
        "fix_the_data", "add_bounds_check", "handle_edge_case",
    ])
    _HASH_BUCKETS = 32  # Fixed number of hash buckets for rare values

    @staticmethod
    def _hash_bucket(value: str, num_buckets: int = 32) -> int:
        """Map a string value to a fixed hash bucket index."""
        import hashlib as _hl
        return int(_hl.md5(value.encode()).hexdigest(), 16) % num_buckets

    def _extract_features(self, signal: LearningSignal) -> Dict[str, float]:
        """Extract rich ML features from signal.

        Includes exception type, strategy, message analysis, code context,
        temporal features, and historical performance.

        F-09: Uses feature hashing for rare exception types and strategies
        to prevent unbounded feature-space growth while preserving
        explicit features for common known values.

        Args:
            signal: Learning signal

        Returns:
            Feature dict for ML model
        """
        features: Dict[str, float] = {}

        # Exception type features -- explicit for known types, hashed for rare
        exc_type = signal.exception_type.lower()
        if exc_type in self._KNOWN_EXCEPTIONS:
            features[f"exc_{exc_type}"] = 1.0
        else:
            bucket = self._hash_bucket(exc_type, self._HASH_BUCKETS)
            features[f"exc_hash_{bucket}"] = 1.0

        # Strategy features -- explicit for known strategies, hashed for rare
        if signal.strategy_used:
            strategy = signal.strategy_used.lower()
            if strategy in self._KNOWN_STRATEGIES:
                features[f"strategy_{signal.strategy_used}"] = 1.0
            else:
                bucket = self._hash_bucket(strategy, self._HASH_BUCKETS)
                features[f"strategy_hash_{bucket}"] = 1.0

        # Message-based features
        message = signal.exception_message.lower()
        features["message_length"] = float(len(signal.exception_message))
        features["has_key_in_message"] = 1.0 if "key" in message else 0.0
        features["has_none_in_message"] = 1.0 if "none" in message else 0.0
        features["has_type_in_message"] = 1.0 if "type" in message else 0.0
        features["has_attribute_in_message"] = 1.0 if "attribute" in message else 0.0
        features["has_index_in_message"] = 1.0 if "index" in message else 0.0
        features["has_import_in_message"] = 1.0 if "import" in message else 0.0
        features["has_module_in_message"] = 1.0 if "module" in message else 0.0

        # Message token count (unique tokens for complexity signal)
        msg_tokens = set(message.split())
        features["message_unique_tokens"] = float(len(msg_tokens))

        # File extension features (from signal.features)
        file_path = signal.features.get("file", "")
        if isinstance(file_path, str) and "." in file_path:
            ext = file_path.rsplit(".", 1)[-1].lower()
            for lang_ext in ("py", "js", "ts", "go", "java"):
                features[f"file_ext_{lang_ext}"] = 1.0 if ext == lang_ext else 0.0

        # Stack depth
        stack_depth = signal.features.get("stack_depth", 0)
        if isinstance(stack_depth, (int, float)):
            features["stack_depth"] = float(stack_depth)

        # Code context features
        code_context = signal.features.get("code_context", "")
        if isinstance(code_context, str) and code_context:
            features["has_async"] = 1.0 if "async" in code_context or "await" in code_context else 0.0
            features["has_mock"] = 1.0 if "mock" in code_context.lower() or "patch" in code_context.lower() else 0.0
            features["has_fixture"] = 1.0 if "fixture" in code_context.lower() or "@pytest" in code_context.lower() else 0.0

        # Diff size
        diff_size = signal.features.get("diff_size", 0)
        if isinstance(diff_size, (int, float)):
            features["diff_size"] = float(diff_size)

        # Temporal feature: hour of day
        features["hour_of_day"] = float(signal.timestamp.hour)

        # Historical success rate for this strategy (from Thompson)
        if self.thompson and signal.strategy_used:
            features["strategy_success_history"] = self.get_expected_success_rate(
                signal.exception_type, signal.strategy_used
            )

        # Performance monitor state
        features["monitor_acceptance_rate"] = self.monitor.rolling_acceptance_rate
        features["in_fallback"] = 1.0 if self.monitor.in_fallback else 0.0

        # Add any custom features from signal
        for key, value in signal.features.items():
            if isinstance(value, (int, float)) and key not in (
                "stack_depth", "diff_size"
            ):
                features[key] = float(value)

        return features

    def _store_positive_example(self, signal: LearningSignal) -> None:
        """Store accepted fix as a positive example.

        Args:
            signal: Learning signal for accepted fix
        """
        from autodebug.redaction import redact_value
        from worker.privacy import redact_pii

        def _redact(text: str) -> str:
            """Apply both secret redaction and PII redaction."""
            return redact_pii(redact_value(text))

        redacted_features = dict(signal.features)
        for key in ("code_context", "file"):
            if key in redacted_features:
                redacted_features[key] = _redact(str(redacted_features[key]))

        example = {
            "exception_type": signal.exception_type,
            "exception_message": _redact(signal.exception_message),
            "strategy": signal.strategy_used,
            "fix_id": signal.fix_id,
            "accepted_at": signal.timestamp.isoformat(),
            "features": redacted_features,
        }
        if signal.user_modifications:
            example["user_modifications"] = _redact(signal.user_modifications)
        self.persistence.store_example(example)

    def _track_rejection(self, signal: LearningSignal) -> None:
        """Track rejection for pattern detection.

        Args:
            signal: Learning signal for rejected fix
        """
        key = f"{signal.exception_type}:{signal.strategy_used}"
        self.persistence.record_rejection(key)

        # Check for repeated rejections
        recent_count = self.persistence.count_recent_rejections(key, days=7)
        if recent_count >= 5:
            from autodebug.redaction import redact_value
            from worker.privacy import redact_pii

            self.persistence.create_alert({
                "type": "repeated_rejection",
                "pattern": key,
                "rejection_count": recent_count,
                "sample_error": redact_pii(redact_value(signal.exception_message[:200])),
                "created_at": datetime.now(timezone.utc).isoformat(),
            })

    def get_strategy_priority(
        self,
        exception_type: str,
        strategy_id: str,
    ) -> float:
        """Get learned priority adjustment using Thompson Sampling.

        Args:
            exception_type: Type of exception
            strategy_id: Strategy identifier

        Returns:
            Sampled priority (0-1)
        """
        if self.thompson is None:
            return 0.5

        stats = self.thompson.get_stats(exception_type, strategy_id)
        return stats.sample()

    def get_expected_success_rate(
        self,
        exception_type: str,
        strategy_id: str,
    ) -> float:
        """Get expected success rate for a strategy.

        Args:
            exception_type: Type of exception
            strategy_id: Strategy identifier

        Returns:
            Expected success rate (0-1)
        """
        if self.thompson is None:
            return 0.5

        stats = self.thompson.get_stats(exception_type, strategy_id)
        return stats.success_rate

    def _compute_uncertainty(self, prediction_prob: float) -> float:
        """Compute uncertainty of a prediction for active sample selection.

        Uses proximity to the decision boundary (0.5) as a measure of
        model uncertainty. Returns 1.0 when the model is maximally
        uncertain (prediction at 0.5) and 0.0 when fully confident
        (prediction at 0.0 or 1.0).

        F-15: When the model has poor calibration (high ECE), uncertainty
        is boosted so that more samples get extra learning passes,
        which helps the model recalibrate faster.

        Args:
            prediction_prob: Model's predicted probability of success (0-1)

        Returns:
            Uncertainty score in [0, 1], where 1.0 = maximum uncertainty
        """
        base_uncertainty = 1.0 - abs(2.0 * prediction_prob - 1.0)

        # F-15: Calibration-aware boost.  When ECE is high (model is
        # miscalibrated), increase the uncertainty score so that more
        # samples trigger extra learning passes, helping the model
        # recalibrate.  This implements the insight from "Calibrated
        # Uncertainty Sampling" (arXiv 2510.03162).
        try:
            ece = self._get_calibration_error()
            if ece > 0.1:  # Miscalibrated: boost uncertainty
                calibration_boost = min(ece, 0.3)  # Cap boost at 0.3
                return min(1.0, base_uncertainty + calibration_boost)
        except Exception:
            pass

        return base_uncertainty

    def _get_calibration_error(self) -> float:
        """Get current Expected Calibration Error from metrics dashboard.

        F-15: Retrieves ECE from CalibrationMetrics if available.

        Returns:
            ECE value (0 = perfectly calibrated), or 0.0 if unavailable
        """
        try:
            from autodebug.learning.metrics_dashboard import CalibrationMetrics

            # Use a shared instance if one exists; otherwise return 0.0
            if hasattr(self, "_calibration_metrics"):
                return self._calibration_metrics.get_ece()
        except ImportError:
            pass
        return 0.0

    def predict_fix_success(self, features: Dict[str, float]) -> float:
        """Predict success probability for a fix.

        Args:
            features: Feature dict for the fix

        Returns:
            Predicted success probability (0-1)
        """
        return self.predictor.predict_success(features)

    def get_best_examples(
        self,
        exception_type: str,
        exception_message: str,
        limit: int = 3,
    ) -> List[Dict[str, Any]]:
        """Get best examples from accepted fixes.

        Args:
            exception_type: Type of exception
            exception_message: Exception message
            limit: Maximum examples to return

        Returns:
            List of example dicts
        """
        # Get examples for this exception type
        examples = self.persistence.get_examples(
            exception_type=exception_type,
            limit=limit * 2,  # Get more, then filter
        )

        # Add success rate from Thompson Sampling
        enriched = []
        for ex in examples:
            strategy = ex.get("strategy", "")
            if strategy:
                success_rate = self.get_expected_success_rate(exception_type, strategy)
                ex["success_rate"] = success_rate
            enriched.append(ex)

        # Sort by success rate and return top N
        enriched.sort(key=lambda x: x.get("success_rate", 0), reverse=True)
        return enriched[:limit]

    def get_stats(self) -> Dict[str, Any]:
        """Get learning statistics.

        Returns:
            Dict with learning stats
        """
        stats = {
            "feedback_counts": self.feedback_counts.copy(),
            "predictor_accuracy": self.predictor.get_accuracy(),
            "predictor_rolling_accuracy": self.predictor.get_rolling_accuracy(),
            "predictor_log_loss": self.predictor.get_log_loss(),
            "predictor_samples": self.predictor.sample_count,
            "predictor_available": True,
            "stored_examples": len(self.persistence.examples),
            "active_alerts": len(self.persistence.alerts),
            "monitor_acceptance_rate": self.monitor.rolling_acceptance_rate,
            "monitor_in_fallback": self.monitor.in_fallback,
            "monitor_fallback_count": self.monitor.fallback_count,
            "monitor_drift_count": self.monitor.drift_count,
            "monitor_drift_detector": self.monitor.drift_detector_name,
            "monitor_last_drift_at": (
                self.monitor.last_drift_at.isoformat()
                if self.monitor.last_drift_at
                else None
            ),
        }

        # Add Thompson stats if available
        if self.thompson:
            all_stats = self.thompson.get_all_stats()
            stats["thompson_strategies"] = len(all_stats)
            if all_stats:
                rates = [s["rate"] for s in all_stats.values()]
                stats["thompson_avg_success_rate"] = sum(rates) / len(rates)

        # Add cross-org stats if available
        if self._cross_org:
            stats["cross_org_enabled"] = True
            stats["cross_org_org_id"] = self._org_id

        return stats

    # --------------------- Cross-Org Integration ---------------------

    @property
    def cross_org(self) -> Optional[CrossOrgAggregator]:
        """Get cross-org aggregator if configured."""
        return self._cross_org

    @property
    def pattern_matcher(self) -> Optional[PatternMatcher]:
        """Get or create pattern matcher for cross-org patterns."""
        if self._cross_org and self._pattern_matcher is None:
            from autodebug.learning.pattern_matcher import PatternMatcher

            self._pattern_matcher = PatternMatcher(self._cross_org)
        return self._pattern_matcher

    def get_global_patterns(
        self,
        exception_type: str,
        exception_message: str,
        stack_trace: List[str],
        max_results: int = 5,
    ) -> List[PatternMatch]:
        """Get matching patterns from cross-organization database.

        Args:
            exception_type: Type of exception
            exception_message: Exception message
            stack_trace: Stack trace lines
            max_results: Maximum patterns to return

        Returns:
            List of matched patterns with scores
        """
        if not self.pattern_matcher:
            return []

        return self.pattern_matcher.find_matching_patterns(
            exception_type=exception_type,
            exception_message=exception_message,
            stack_trace=stack_trace,
            max_results=max_results,
        )

    def get_global_strategy_recommendation(
        self,
        exception_type: str,
        exception_message: str,
        stack_trace: List[str],
    ) -> Optional[str]:
        """Get best strategy recommendation from global patterns.

        Args:
            exception_type: Type of exception
            exception_message: Exception message
            stack_trace: Stack trace lines

        Returns:
            Recommended strategy name or None
        """
        if not self.pattern_matcher:
            return None

        return self.pattern_matcher.get_best_strategy(
            exception_type=exception_type,
            exception_message=exception_message,
            stack_trace=stack_trace,
        )

    def should_use_global_patterns(
        self,
        exception_type: str,
        exception_message: str,
    ) -> bool:
        """Determine if global patterns should be consulted.

        Uses global patterns when:
        - Few local examples exist
        - Exception is commonly seen across organizations

        Args:
            exception_type: Type of exception
            exception_message: Exception message

        Returns:
            True if global patterns should be used
        """
        if not self._cross_org:
            return False

        # Check local example count
        local_examples = self.persistence.get_examples(
            exception_type=exception_type,
            limit=10,
        )

        # Use global if few local examples
        if len(local_examples) < 5:
            return True

        # Check if pattern is common across orgs
        if self.pattern_matcher:
            from autodebug.learning.pattern_matcher import PatternRecommendationEngine

            engine = PatternRecommendationEngine(self.pattern_matcher)
            return engine.should_use_global_patterns(
                exception_type=exception_type,
                exception_message=exception_message,
                local_history_count=len(local_examples),
            )

        return False

    def contribute_to_global(
        self,
        exception_type: str,
        exception_message: str,
        stack_trace: List[str],
        strategy_used: str,
        was_accepted: bool,
    ) -> bool:
        """Contribute a fix pattern to cross-org learning.

        Anonymizes and contributes successful fix patterns to the
        global database for other organizations to learn from.

        Args:
            exception_type: Type of exception
            exception_message: Exception message
            stack_trace: Stack trace lines
            strategy_used: Strategy that was used
            was_accepted: Whether the fix was accepted

        Returns:
            True if contribution was successful
        """
        if not self._cross_org or not self._org_id:
            return False

        try:
            self._cross_org.contribute_pattern(
                org_id=self._org_id,
                exception_type=exception_type,
                exception_message=exception_message,
                stack_trace=stack_trace,
                fix_strategy=strategy_used,
                was_accepted=was_accepted,
            )
            return True
        except Exception as e:
            logger.debug(f"Failed to contribute to global patterns: {e}")
            return False

    def get_combined_strategy_rankings(
        self,
        exception_type: str,
        exception_message: str,
        stack_trace: List[str],
    ) -> List[tuple[str, float, str]]:
        """Get strategy rankings combining local and global learning.

        Combines Thompson Sampling local learning with cross-org
        pattern matching for comprehensive strategy recommendations.

        Args:
            exception_type: Type of exception
            exception_message: Exception message
            stack_trace: Stack trace lines

        Returns:
            List of (strategy, score, source) tuples
        """
        rankings: Dict[str, tuple[float, str]] = {}

        # Get local Thompson Sampling recommendations
        if self.thompson:
            # Get all strategies for this exception type
            all_stats = self.thompson.get_all_stats()
            for key, stats in all_stats.items():
                if key.startswith(f"{exception_type}:"):
                    strategy = key.split(":", 1)[1]
                    rankings[strategy] = (stats["rate"], "local")

        # Get global pattern recommendations
        if self.pattern_matcher:
            global_rankings = self.pattern_matcher.get_strategy_rankings(
                exception_type=exception_type,
                exception_message=exception_message,
                stack_trace=stack_trace,
            )

            for strategy, score in global_rankings:
                if strategy in rankings:
                    # Blend local and global scores
                    local_score, _ = rankings[strategy]
                    blended = 0.7 * local_score + 0.3 * score
                    rankings[strategy] = (blended, "blended")
                else:
                    rankings[strategy] = (score, "global")

        # Convert to sorted list
        result = [
            (strategy, score, source)
            for strategy, (score, source) in rankings.items()
        ]
        result.sort(key=lambda x: x[1], reverse=True)

        return result

    # --------------------- Checkpoint / Restore ---------------------

    # F-10: Model versioning with rollback support.
    # Maintains a history of recent checkpoints so the model can be
    # rolled back to a previous known-good version if quality degrades.
    _MAX_CHECKPOINT_HISTORY = 5

    def _store_checkpoint_history(self, checkpoint_data: Dict[str, Any]) -> None:
        """Store checkpoint in the rolling history for rollback support.

        Args:
            checkpoint_data: Checkpoint to archive
        """
        if not hasattr(self, "_checkpoint_history"):
            self._checkpoint_history: List[Dict[str, Any]] = []

        self._checkpoint_history.append(checkpoint_data)
        # Keep only the most recent checkpoints
        if len(self._checkpoint_history) > self._MAX_CHECKPOINT_HISTORY:
            self._checkpoint_history = self._checkpoint_history[
                -self._MAX_CHECKPOINT_HISTORY:
            ]

    async def rollback_to_previous(self) -> Dict[str, Any]:
        """Roll back to the previous checkpoint.

        F-10: When model quality degrades (e.g., PerformanceMonitor
        triggers fallback), this method restores the second-most-recent
        checkpoint instead of continuing with the current degraded model.

        Returns:
            Dict with rollback status
        """
        if not hasattr(self, "_checkpoint_history") or len(self._checkpoint_history) < 2:
            return {
                "status": "error",
                "message": "No previous checkpoint available for rollback",
            }

        # Remove the current (bad) checkpoint
        self._checkpoint_history.pop()
        # Restore the previous one
        previous = self._checkpoint_history[-1]
        result = await self.restore(previous)
        result["rollback"] = True
        logger.info(
            "Model rolled back to checkpoint from %s",
            previous.get("timestamp", "unknown"),
        )
        return result

    @staticmethod
    def _get_hmac_key(org_id: Optional[str]) -> bytes:
        """Derive HMAC key for pickle signing/verification.

        Uses a proper secret (SECRET_KEY env var) rather than the
        non-secret org_id alone. Falls back to org_id combined with
        a machine-specific value if no secret is configured.

        Args:
            org_id: Organization identifier (used as salt, not sole key)

        Returns:
            32-byte HMAC key
        """
        import hashlib
        import os
        import uuid

        # Prefer a dedicated secret from the environment
        secret = os.environ.get("SECRET_KEY") or os.environ.get("SIGNING_KEY") or ""

        if not secret:
            # Fallback: combine org_id with machine-specific node ID.
            # uuid.getnode() returns a hardware-based identifier that
            # is not guessable from the org_id alone.
            machine_id = str(uuid.getnode())
            secret = f"hikaflow-fallback-{machine_id}"

        key_material = f"{secret}:hikaflow-learner:{org_id or 'default'}"
        return hashlib.sha256(key_material.encode()).digest()

    async def checkpoint(self) -> Dict[str, Any]:
        """Save current learning state to persistent storage.

        Saves:
        - Thompson Sampling state (success/failure counts per strategy)
        - Online learner model state
        - Feedback counts
        - Persistence state (examples, rejections, alerts)

        Returns:
            Dict with checkpoint status and metadata.
        """
        checkpoint_data: Dict[str, Any] = {
            "version": "2.0",
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "org_id": self._org_id,
            "feedback_counts": self.feedback_counts.copy(),
        }

        # Checkpoint Thompson Sampling state
        if self.thompson:
            try:
                all_stats = self.thompson.get_all_stats()
                thompson_state = {}
                for key, stats in all_stats.items():
                    thompson_state[key] = {
                        "successes": stats.get("successes", 0),
                        "failures": stats.get("failures", 0),
                        "rate": stats.get("rate", 0.5),
                    }
                checkpoint_data["thompson_state"] = thompson_state
            except Exception as e:
                logger.warning(f"Failed to checkpoint Thompson state: {e}")
                checkpoint_data["thompson_error"] = str(e)

        # Checkpoint online learner (River model) with HMAC signing
        try:
            import base64
            import hashlib
            import hmac
            import pickle

            model_bytes = pickle.dumps(self.predictor._model)
            checkpoint_data["predictor_model"] = base64.b64encode(model_bytes).decode("utf-8")
            checkpoint_data["predictor_samples"] = self.predictor.sample_count
            # Sign with HMAC using proper secret key (not just org_id)
            mac_key = self._get_hmac_key(self._org_id)
            checkpoint_data["predictor_model_mac"] = hmac.new(
                mac_key, model_bytes, hashlib.sha256
            ).hexdigest()
        except Exception as e:
            logger.warning(f"Failed to checkpoint predictor: {e}")
            checkpoint_data["predictor_error"] = str(e)

        # Checkpoint persistence state
        try:
            checkpoint_data["persistence"] = {
                "examples_count": len(self.persistence.examples),
                "rejections_count": len(self.persistence.rejections),
                "alerts_count": len(self.persistence.alerts),
            }

            # Save to persistence path if configured
            if self.persistence.path:
                self.persistence._save()
        except Exception as e:
            logger.warning(f"Failed to checkpoint persistence: {e}")
            checkpoint_data["persistence_error"] = str(e)

        # F-10: Store in local history for rollback support
        self._store_checkpoint_history(checkpoint_data)

        # Store checkpoint in database if org_id is set
        if self._org_id:
            try:
                from api import db
                db.store_learning_checkpoint(self._org_id, checkpoint_data)
            except Exception as e:
                logger.warning(f"Failed to store checkpoint in database: {e}")

        return {
            "status": "success",
            "checkpoint_data": checkpoint_data,
        }

    async def restore(self, checkpoint_data: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """Restore learning state from checkpoint.

        Args:
            checkpoint_data: Checkpoint data to restore from.
                            If None, loads from database.

        Returns:
            Dict with restore status.
        """
        if checkpoint_data is None and self._org_id:
            # Load from database
            try:
                from api import db
                checkpoint_data = db.get_learning_checkpoint(self._org_id)
            except Exception as e:
                return {
                    "status": "error",
                    "message": f"Failed to load checkpoint: {e}",
                }

        if not checkpoint_data:
            return {
                "status": "skipped",
                "message": "No checkpoint data available",
            }

        restored = []

        # Restore feedback counts
        if "feedback_counts" in checkpoint_data:
            self.feedback_counts = checkpoint_data["feedback_counts"].copy()
            restored.append("feedback_counts")

        # Restore Thompson Sampling state
        if "thompson_state" in checkpoint_data and self.thompson:
            try:
                for key, state in checkpoint_data["thompson_state"].items():
                    # Parse key as "exception_type:strategy"
                    if ":" in key:
                        exception_type, strategy = key.split(":", 1)
                        # Directly set Beta parameters instead of replaying events
                        # (avoids O(n) replay and incorrect temporal decay)
                        successes = state.get("successes", 0)
                        failures = state.get("failures", 0)
                        if hasattr(self.thompson, 'set_stats'):
                            self.thompson.set_stats(
                                exception_type, strategy,
                                successes=successes, failures=failures,
                            )
                        else:
                            # Fallback: replay if set_stats not available
                            for _ in range(successes):
                                self.thompson.update(exception_type, strategy, success=True)
                            for _ in range(failures):
                                self.thompson.update(exception_type, strategy, success=False)
                restored.append("thompson_state")
            except Exception as e:
                logger.warning(f"Failed to restore Thompson state: {e}")

        # Restore predictor model (using HMAC-verified pickle for safety)
        if "predictor_model" in checkpoint_data:
            try:
                import base64
                import hashlib
                import hmac
                import pickle

                model_bytes = base64.b64decode(checkpoint_data["predictor_model"])
                if len(model_bytes) > 5_000_000:
                    logger.warning("Predictor model too large, skipping restore")
                    model_bytes = b""
                stored_mac = checkpoint_data.get("predictor_model_mac", "")
                # Verify HMAC before deserializing (key from proper secret)
                mac_key = self._get_hmac_key(self._org_id)
                expected_mac = hmac.new(mac_key, model_bytes, hashlib.sha256).hexdigest()
                if not hmac.compare_digest(stored_mac, expected_mac):
                    logger.warning("Predictor model HMAC verification failed, skipping restore")
                else:
                    # Safe path: only restore if HMAC is valid and size-bounded
                    self.predictor._model = pickle.loads(model_bytes)  # noqa: S301  # nosemgrep: pickle-untrusted-data
                    self.predictor._sample_count = checkpoint_data.get("predictor_samples", 0)
                    restored.append("predictor_model")
            except Exception as e:
                logger.warning(f"Failed to restore predictor: {e}")

        return {
            "status": "success",
            "restored": restored,
            "checkpoint_timestamp": checkpoint_data.get("timestamp"),
        }
