"""
Active Inference Engine v2 — The Bot's Immune System.

Bayesian Active Inference: predicts future UI states before acting,
evaluates predictions against reality, and steers behavior based on
accumulated surprise (Free Energy).

v2 Enhancements:
- Consecutive error tracking → automatic policy escalation
- Interaction throttling → reduces follow/like probability under high surprise
- Session abort recommendation → when environment is fundamentally unstable
- Rich prediction context → tracks WHAT was expected vs. WHAT was found
"""

import logging
import math
import time

from colorama import Fore

logger = logging.getLogger(__name__)


class ActiveInferenceEngine:
    """
    Bayesian Active Inference Engine.
    Calculates Free Energy (Surprise) based on prediction errors in the
    Instagram environment. Steers the agent's 'Thermodynamic Policy'.

    Policies:
    - STABLE:   Free energy < 0.75. Normal operation. All interactions enabled.
    - CAUTIOUS: Free energy 0.75-1.2. Reduced interaction probability. Longer waits.
    - DORMANT:  Free energy > 1.2. Minimal interactions. Maximum sleep. May recommend abort.
    """

    def __init__(self, username):
        self.username = username
        self.free_energy = 0.0
        self.surprise_threshold = 0.75
        self.last_update = time.time()
        self.policy = "STABLE"  # STABLE, CAUTIOUS, DORMANT
        self.expectation_history = []

        # v2: Consecutive error tracking for escalation
        self._consecutive_prediction_errors = 0
        self._total_predictions = 0
        self._total_errors = 0
        self._session_start = time.time()

    def calculate_surprise(self, predicted_outcome: float, observed_outcome: float):
        """
        Bayesian surprise calculation (simplified Kullback-Leibler divergence).
        """
        # prediction error
        error = abs(predicted_outcome - observed_outcome)

        # Free energy accumulation
        self.free_energy = (self.free_energy * 0.7) + (error * 0.3)

        # Decay free energy over time (Thermodynamic relaxation)
        now = time.time()
        hours_passed = (now - self.last_update) / 3600.0
        decay = math.exp(-0.1 * hours_passed)
        self.free_energy *= decay
        self.last_update = now

        # Policy steering
        if self.free_energy > 1.2:
            self.policy = "DORMANT"
        elif self.free_energy > self.surprise_threshold:
            self.policy = "CAUTIOUS"
        else:
            self.policy = "STABLE"

        logger.info(
            f"⚖️ [Active Inference] Surprise: {self.free_energy:.4f} | Policy: {self.policy}",
            extra={"color": f"{Fore.BLUE}"},
        )
        return self.free_energy

    def predict_state(self, expected_signature: list):
        """
        Registers an expectation about the future UI state before acting.
        expected_signature: list of terms expected in the resulting XML.
        """
        self.expectation_history.append(expected_signature)
        logger.debug(
            f"⚖️ [Shadow Mode] Predicting future state containing: {expected_signature}", extra={"color": f"{Fore.BLUE}"}
        )

    def evaluate_prediction(self, context_xml: str) -> bool:
        """
        Evaluates the last prediction against reality.
        Returns True if reality matches prediction, False otherwise (Prediction Error).

        v2: Tracks consecutive errors and escalates policy automatically.
        """
        if not self.expectation_history:
            return True

        expected_signature = self.expectation_history.pop()
        self._total_predictions += 1
        matched = any(sig.lower() in context_xml.lower() for sig in expected_signature)

        if matched:
            self._consecutive_prediction_errors = 0
            self.calculate_surprise(1.0, 1.0)
            return True
        else:
            self._consecutive_prediction_errors += 1
            self._total_errors += 1
            logger.warning(
                f"⚖️ [Shadow Mode] Prediction Error #{self._consecutive_prediction_errors}! "
                f"Did not find {expected_signature} in resulting UI.",
                extra={"color": f"{Fore.RED}"},
            )
            self.calculate_surprise(1.0, 0.0)

            # v2: Consecutive error escalation
            if self._consecutive_prediction_errors >= 5:
                self.policy = "DORMANT"
                logger.error(
                    f"🚨 [Active Inference] {self._consecutive_prediction_errors} consecutive prediction errors! "
                    f"Environment is fundamentally unstable. DORMANT mode engaged.",
                    extra={"color": f"{Fore.RED}"},
                )
            elif self._consecutive_prediction_errors >= 3:
                self.policy = "CAUTIOUS"
                logger.warning(
                    f"⚠️ [Active Inference] {self._consecutive_prediction_errors} consecutive errors. "
                    f"Switching to CAUTIOUS policy.",
                    extra={"color": f"{Fore.YELLOW}"},
                )

            # ── Dojo Data Engine Hook ──
            # When prediction fails, explicitly submit the snapshot for shadow-compilation
            try:
                from GramAddict.core.dojo_engine import DojoEngine

                # Note: get_instance() works without passing device as it was already initialized in bot_flow by this point.
                dojo = DojoEngine.get_instance()
                dojo.submit_snapshot(
                    heuristic_name=str(expected_signature),
                    context_xml=context_xml,
                    intent_prompt=f"Locate the missing elements or correct the heuristic predicting state: {expected_signature}",
                )
            except Exception as e:
                logger.error(f"Failed to offload snapshot to Dojo Engine: {e}")

            return False

    def get_sleep_modifier(self):
        """
        Returns a multiplier for sleep durations based on surprise.
        """
        if self.policy == "DORMANT":
            return 5.0
        if self.policy == "CAUTIOUS":
            return 2.0
        return 1.0

    # ──────────────────────────────────────────────
    # v2: New behavioral steering methods
    # ──────────────────────────────────────────────

    def get_interaction_probability(self) -> float:
        """
        Returns a probability multiplier [0.0 - 1.0] for interaction decisions.

        Under STABLE: 1.0 (full interaction rate)
        Under CAUTIOUS: 0.5 (halved interaction rate)
        Under DORMANT: 0.1 (minimal interaction — only high-confidence targets)

        This directly modifies follow/like/comment probability in the feed loop.
        """
        if self.policy == "DORMANT":
            return 0.1
        if self.policy == "CAUTIOUS":
            return 0.5
        return 1.0

    def should_abort_session(self) -> bool:
        """
        Recommends session abort when the environment is fundamentally broken.

        Triggers:
        - 5+ consecutive prediction errors (UI is completely unexpected)
        - Free energy > 2.0 (accumulated instability beyond recovery)

        The caller (bot_flow) can choose to honor this or override.
        """
        if self._consecutive_prediction_errors >= 5:
            return True
        if self.free_energy > 2.0:
            return True
        return False

    def get_error_rate(self) -> float:
        """Returns the session-wide prediction error rate."""
        if self._total_predictions == 0:
            return 0.0
        return self._total_errors / self._total_predictions

    def get_diagnostics(self) -> dict:
        """Returns a diagnostic snapshot for logging/telemetry."""
        return {
            "free_energy": round(self.free_energy, 4),
            "policy": self.policy,
            "consecutive_errors": self._consecutive_prediction_errors,
            "total_predictions": self._total_predictions,
            "total_errors": self._total_errors,
            "error_rate": round(self.get_error_rate(), 4),
            "session_uptime_minutes": round((time.time() - self._session_start) / 60, 1),
            "should_abort": self.should_abort_session(),
        }
