"""
InsAIts SDK - Persistence Engine (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__)

# ---------------------------------------------------------------------------
# Public constants
# ---------------------------------------------------------------------------

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

MAX_LEVEL = LEVEL_CIRCUIT_BREAK

RESET_TIMEOUT_SECONDS = 600  # 10 minutes
TOOLS_TRIED_BACKOFF_THRESHOLD = 3
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_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
# ---------------------------------------------------------------------------

@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_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.
    """

    def __init__(self, stealth: bool = False) -> None:
        self._stealth = stealth
        self._states: Dict[Tuple[str, str], _AgentTaskState] = {}

    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."""
        key = (agent_id, task_context)
        state = self._get_or_create_state(key)
        now = time.monotonic()

        # Check timeout reset
        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

        # Reset tools-since-last-surrender for next cycle
        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)

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

    def get_escalation_level(self, agent_id: str, task_context: str) -> int:
        """Return the current escalation level (0 if no failures)."""
        key = (agent_id, task_context)
        state = self._states.get(key)
        if state is None:
            return 0

        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."""
        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."""
        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."""
        key = (agent_id, task_context)
        state = self._states.get(key)
        if state is None:
            return False

        if len(state.tools_since_last_surrender) >= TOOLS_TRIED_BACKOFF_THRESHOLD:
            return True

        return False

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

    def _get_or_create_state(
        self, key: Tuple[str, str]
    ) -> _AgentTaskState:
        if key not in self._states:
            self._states[key] = _AgentTaskState()
        return self._states[key]

    @staticmethod
    def _compute_raw_level(failure_count: int) -> int:
        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:
        level = raw_level

        # Rule 1: tools diversity -> reset to level 1
        if len(state.tools_since_last_surrender) >= TOOLS_TRIED_BACKOFF_THRESHOLD:
            return LEVEL_NUDGE

        # Rule 2: different error type -> back off 1 level
        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)

        return level

    @staticmethod
    def _should_timeout_reset(state: _AgentTaskState, now: float) -> bool:
        if state.last_failure_ts == 0.0:
            return False
        return (now - state.last_failure_ts) >= RESET_TIMEOUT_SECONDS

    @staticmethod
    def _reset_state(state: _AgentTaskState) -> None:
        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:
        if not self._stealth:
            return True
        return level >= STEALTH_VISIBLE_THRESHOLD


__all__ = [
    "LEVEL_NUDGE",
    "LEVEL_REDIRECT",
    "LEVEL_DECOMPOSE",
    "LEVEL_ESCALATE",
    "LEVEL_CIRCUIT_BREAK",
    "MAX_LEVEL",
    "RESET_TIMEOUT_SECONDS",
    "TOOLS_TRIED_BACKOFF_THRESHOLD",
    "STEALTH_VISIBLE_THRESHOLD",
    "EscalationAction",
    "PersistenceEngine",
]
