# Copyright (c) 2024-2026 YuyAI / InsAIts Team. All rights reserved.
# Proprietary and confidential. See LICENSE.premium for terms.
"""
InsAIts SDK - Persistence Engine - Premium Core (Phase 3.6-3.11)
=================================================================
Intelligent escalation when agents surrender, loop, or go NPC.

Key principle: pressure only when agent is NOT genuinely stuck.
If agent is trying new approaches (different errors, different tools),
back off. If agent is repeating the same thing or giving up, escalate.

Escalation levels:
    1 - NUDGE:         1st surrender/NPC
    2 - REDIRECT:      2nd surrender within 5 min
    3 - DECOMPOSE:     3rd surrender or loop detected
    4 - ESCALATE:      4+ surrenders, same task
    5 - CIRCUIT_BREAK: 5+ surrenders, no progress

Author: InsAIts SDK
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass
from typing import Dict, Optional, Set, Tuple

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

LEVEL_NUDGE = 1
LEVEL_REDIRECT = 2
LEVEL_DECOMPOSE = 3
LEVEL_ESCALATE = 4
LEVEL_CIRCUIT_BREAK = 5

MAX_LEVEL = LEVEL_CIRCUIT_BREAK

# After this many seconds with no failures, escalation resets.
RESET_TIMEOUT_SECONDS = 600  # 10 minutes

# If an agent has tried this many distinct tools since last surrender,
# it is actively exploring and we back off to level 1.
TOOLS_TRIED_BACKOFF_THRESHOLD = 3

# Stealth: levels 1-3 are silent (audit-only), 4-5 are always visible.
STEALTH_VISIBLE_THRESHOLD = LEVEL_ESCALATE  # level 4+

# ---------------------------------------------------------------------------
# Intervention messages per level
# ---------------------------------------------------------------------------

_INTERVENTION_MESSAGES: Dict[int, str] = {
    LEVEL_NUDGE: (
        "Try a different approach. Have you considered breaking the "
        "problem down or searching for similar patterns in the codebase?"
    ),
    LEVEL_REDIRECT: (
        "Break this into smaller steps. What is the specific blocker? "
        "Isolate the exact failure point before attempting a fix."
    ),
    LEVEL_DECOMPOSE: (
        "List 3 alternative approaches to solve this. "
        "Try the simplest one first. Do not repeat what already failed."
    ),
    LEVEL_ESCALATE: (
        "This task may need human input. Summarize what you have tried, "
        "what worked, what failed, and what is blocking you."
    ),
    LEVEL_CIRCUIT_BREAK: (
        "STOP. You are stuck in a loop. Ask the user for guidance "
        "before continuing. Do not attempt another solution without "
        "explicit direction."
    ),
}

# Context-aware suffixes appended when task_context is available.
_CONTEXT_SUFFIX: Dict[int, str] = {
    LEVEL_NUDGE: " Consider approaching '{task}' from a different angle.",
    LEVEL_REDIRECT: " Focus specifically on the failing step of '{task}'.",
    LEVEL_DECOMPOSE: " For '{task}', what are the atomic sub-problems?",
    LEVEL_ESCALATE: " Regarding '{task}': present a clear status report.",
    LEVEL_CIRCUIT_BREAK: " Regarding '{task}': stop and wait for user input.",
}


# ---------------------------------------------------------------------------
# EscalationAction dataclass
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class EscalationAction:
    """Result of a persistence engine escalation decision.

    Attributes:
        level: Escalation level 1-5.
        message: Intervention text for the agent.
        should_output: Whether to write to stdout (stealth-aware).
        should_log: Whether to write to the audit log (always True).
        should_circuit_break: True only at level 5.
    """

    level: int
    message: str
    should_output: bool
    should_log: bool
    should_circuit_break: bool


# ---------------------------------------------------------------------------
# Per-agent per-task tracking state (ephemeral, session-scoped)
# ---------------------------------------------------------------------------

@dataclass
class _AgentTaskState:
    """Mutable tracking state for one (agent_id, task_context) pair."""

    failure_count: int = 0
    escalation_level: int = 0
    last_error_type: Optional[str] = None
    tools_tried: Set[str] = None  # type: ignore[assignment]
    first_failure_ts: float = 0.0
    last_failure_ts: float = 0.0
    # Tools tried since the last surrender (for back-off calculation)
    tools_since_last_surrender: Set[str] = None  # type: ignore[assignment]

    def __post_init__(self) -> None:
        if self.tools_tried is None:
            self.tools_tried = set()
        if self.tools_since_last_surrender is None:
            self.tools_since_last_surrender = set()


# ---------------------------------------------------------------------------
# PersistenceEngine
# ---------------------------------------------------------------------------

class PersistenceEngine:
    """Intelligent escalation when agents surrender, loop, or go NPC.

    Key principle: pressure only when agent is NOT genuinely stuck.
    If agent is trying new approaches (different errors, different tools),
    back off. If agent is repeating the same thing or giving up, escalate.

    Usage:
        engine = PersistenceEngine(stealth=True)

        action = engine.record_failure(
            agent_id="claude:Bash",
            task_context="fix-login-bug",
            error_type="TypeError",
            tool_name="Bash",
        )
        if action.should_output:
            print(action.message)  # visible only at levels 4-5 in stealth
    """

    def __init__(self, stealth: bool = False) -> None:
        """Initialize the persistence engine.

        Args:
            stealth: If True, levels 1-3 produce audit-only entries
                     (should_output=False). Levels 4-5 are always visible.
        """
        self._stealth = stealth
        # Keyed by (agent_id, task_context)
        self._states: Dict[Tuple[str, str], _AgentTaskState] = {}

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def record_failure(
        self,
        agent_id: str,
        task_context: str,
        error_type: Optional[str] = None,
        tool_name: Optional[str] = None,
    ) -> EscalationAction:
        """Record a failure event and return the appropriate escalation.

        This is the main entry point. Each call increments the failure
        count for the (agent_id, task_context) pair, applies back-off
        logic, computes the new escalation level, and returns an action.

        Args:
            agent_id: Identifier for the agent (e.g. "claude:Bash").
            task_context: Short description of the current task/goal.
            error_type: The type/category of error encountered (optional).
            tool_name: The tool that was being used (optional).

        Returns:
            EscalationAction with level, message, and output flags.
        """
        key = (agent_id, task_context)
        state = self._get_or_create_state(key)
        now = time.monotonic()

        # Check timeout reset: if enough time passed with no failures,
        # reset the state entirely.
        if state.failure_count > 0 and self._should_timeout_reset(state, now):
            self._reset_state(state)

        # Record the failure
        state.failure_count += 1
        if state.first_failure_ts == 0.0:
            state.first_failure_ts = now
        state.last_failure_ts = now

        # Track tools
        if tool_name:
            state.tools_tried.add(tool_name)
            state.tools_since_last_surrender.add(tool_name)

        # Compute raw escalation level from failure count
        raw_level = self._compute_raw_level(state.failure_count)

        # Apply back-off logic
        backed_off = self._apply_backoff(state, error_type, raw_level)

        # Clamp to [1, MAX_LEVEL]
        final_level = max(1, min(backed_off, MAX_LEVEL))

        # Update state
        state.escalation_level = final_level
        state.last_error_type = error_type

        # After recording, reset the tools-since-last-surrender tracker
        # so the next cycle starts fresh
        state.tools_since_last_surrender = set()

        # Build the action
        message = self.get_intervention_message(final_level, task_context)
        should_output = self._should_output(final_level)

        action = EscalationAction(
            level=final_level,
            message=message,
            should_output=should_output,
            should_log=True,
            should_circuit_break=(final_level >= LEVEL_CIRCUIT_BREAK),
        )

        logger.info(
            "PersistenceEngine: agent=%s task=%s failures=%d level=%d backoff=%s",
            agent_id,
            task_context,
            state.failure_count,
            final_level,
            backed_off < raw_level,
        )

        return action

    def get_escalation_level(self, agent_id: str, task_context: str) -> int:
        """Return the current escalation level for an agent/task pair.

        Returns 0 if no failures have been recorded.

        Args:
            agent_id: Agent identifier.
            task_context: Task/goal description.

        Returns:
            Integer escalation level (0 = no failures, 1-5 = active).
        """
        key = (agent_id, task_context)
        state = self._states.get(key)
        if state is None:
            return 0

        # Check timeout
        now = time.monotonic()
        if state.failure_count > 0 and self._should_timeout_reset(state, now):
            self._reset_state(state)
            return 0

        return state.escalation_level

    def get_intervention_message(
        self, level: int, task_context: Optional[str] = None
    ) -> str:
        """Return the intervention message for a given escalation level.

        Args:
            level: Escalation level (1-5).
            task_context: Optional task description for context-aware messages.

        Returns:
            The intervention message string.
        """
        clamped = max(1, min(level, MAX_LEVEL))
        base = _INTERVENTION_MESSAGES[clamped]

        if task_context:
            suffix_template = _CONTEXT_SUFFIX.get(clamped, "")
            if suffix_template:
                base += suffix_template.format(task=task_context[:80])

        return base

    def reset(
        self, agent_id: str, task_context: Optional[str] = None
    ) -> None:
        """Reset tracking for an agent, optionally scoped to a task.

        Args:
            agent_id: Agent identifier.
            task_context: If provided, only reset this specific task.
                          If None, reset all tasks for this agent.
        """
        if task_context is not None:
            key = (agent_id, task_context)
            if key in self._states:
                self._reset_state(self._states[key])
        else:
            keys_to_reset = [
                k for k in self._states if k[0] == agent_id
            ]
            for key in keys_to_reset:
                self._reset_state(self._states[key])

    def should_back_off(self, agent_id: str, task_context: str) -> bool:
        """Check if the agent is making progress and escalation should ease.

        Returns True if the agent has tried different tools or encountered
        different errors, suggesting genuine exploration rather than looping.

        Args:
            agent_id: Agent identifier.
            task_context: Task/goal description.

        Returns:
            True if back-off conditions are met.
        """
        key = (agent_id, task_context)
        state = self._states.get(key)
        if state is None:
            return False

        # Condition 1: agent tried 3+ unique tools since last surrender
        if len(state.tools_since_last_surrender) >= TOOLS_TRIED_BACKOFF_THRESHOLD:
            return True

        # Condition 2: current state reflects a recent error type change
        # (handled during record_failure, but can be queried here too)
        # We expose it as: if last error differs from what it was before
        # the most recent failure, the agent IS trying new things.
        # This is a heuristic -- we can't look back further without history.
        return False

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _get_or_create_state(
        self, key: Tuple[str, str]
    ) -> _AgentTaskState:
        """Get or create tracking state for an (agent_id, task_context) pair."""
        if key not in self._states:
            self._states[key] = _AgentTaskState()
        return self._states[key]

    @staticmethod
    def _compute_raw_level(failure_count: int) -> int:
        """Map failure count to raw escalation level.

        1 failure  -> level 1 (NUDGE)
        2 failures -> level 2 (REDIRECT)
        3 failures -> level 3 (DECOMPOSE)
        4 failures -> level 4 (ESCALATE)
        5+ failures -> level 5 (CIRCUIT_BREAK)
        """
        if failure_count <= 0:
            return 0
        return min(failure_count, MAX_LEVEL)

    def _apply_backoff(
        self,
        state: _AgentTaskState,
        error_type: Optional[str],
        raw_level: int,
    ) -> int:
        """Apply back-off rules to reduce escalation when agent is exploring.

        Rules:
        - If error_type differs from last_error_type -> back off 1 level
        - If tools_since_last_surrender has 3+ unique tools -> back off to 1
        - These stack: tool diversity takes priority (resets to 1).

        Args:
            state: Current agent/task tracking state.
            error_type: The current error type (may be None).
            raw_level: The raw level from failure count alone.

        Returns:
            Adjusted escalation level after back-off.
        """
        level = raw_level

        # Rule 1: tools diversity -- agent is actively exploring
        if len(state.tools_since_last_surrender) >= TOOLS_TRIED_BACKOFF_THRESHOLD:
            level = LEVEL_NUDGE
            logger.debug(
                "Back-off: %d unique tools tried -> reset to level 1",
                len(state.tools_since_last_surrender),
            )
            return level

        # Rule 2: different error type -- agent is trying new approaches
        if (
            error_type is not None
            and state.last_error_type is not None
            and error_type != state.last_error_type
        ):
            level = max(level - 1, LEVEL_NUDGE)
            logger.debug(
                "Back-off: error changed %s -> %s, level %d -> %d",
                state.last_error_type,
                error_type,
                raw_level,
                level,
            )

        return level

    @staticmethod
    def _should_timeout_reset(state: _AgentTaskState, now: float) -> bool:
        """Check if enough time has passed since the last failure to reset."""
        if state.last_failure_ts == 0.0:
            return False
        elapsed = now - state.last_failure_ts
        return elapsed >= RESET_TIMEOUT_SECONDS

    @staticmethod
    def _reset_state(state: _AgentTaskState) -> None:
        """Reset a state object to its initial values."""
        state.failure_count = 0
        state.escalation_level = 0
        state.last_error_type = None
        state.tools_tried = set()
        state.tools_since_last_surrender = set()
        state.first_failure_ts = 0.0
        state.last_failure_ts = 0.0

    def _should_output(self, level: int) -> bool:
        """Determine whether the intervention should produce stdout output.

        In stealth mode, levels 1-3 are silent (audit log only).
        Levels 4-5 are always visible regardless of stealth.
        In non-stealth mode, all levels produce output.
        """
        if not self._stealth:
            return True
        return level >= STEALTH_VISIBLE_THRESHOLD
