"""Round guardrails: tool budgets, retry limits, and failure classification.

Prevents runaway tool loops, classifies failures for structured recovery,
and enforces per-tool-type call budgets within a single agent round.
"""

from __future__ import annotations

from collections import Counter
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

from loguru import logger


# ---------------------------------------------------------------------------
# Failure classification
# ---------------------------------------------------------------------------

class FailureType(str, Enum):
    """Standardized failure categories for tool calls."""

    TIMEOUT = "timeout"
    PERMISSION = "permission"
    EMPTY_RESULT = "empty_result"
    CONFLICT = "conflict"
    INVALID_PARAMS = "invalid_params"
    RUNTIME_ERROR = "runtime_error"
    UNKNOWN = "unknown"

    @classmethod
    def classify(cls, error_text: str) -> FailureType:
        """Classify an error string into a failure type."""
        low = error_text.lower()
        if any(k in low for k in ("timeout", "timed out", "deadline")):
            return cls.TIMEOUT
        if any(k in low for k in ("permission", "denied", "forbidden", "not allowed", "unauthorized")):
            return cls.PERMISSION
        if any(k in low for k in ("not found", "no such", "empty", "no results", "no matching")):
            return cls.EMPTY_RESULT
        if any(k in low for k in ("conflict", "already exists", "locked", "busy")):
            return cls.CONFLICT
        if any(k in low for k in ("invalid", "missing", "required parameter", "bad request")):
            return cls.INVALID_PARAMS
        return cls.UNKNOWN


# ---------------------------------------------------------------------------
# Failure record
# ---------------------------------------------------------------------------

@dataclass
class ToolFailure:
    """A single tool call failure record."""

    tool_name: str
    failure_type: FailureType
    error_text: str
    args_signature: str  # simplified args fingerprint for dedup


# ---------------------------------------------------------------------------
# Round guardrails tracker
# ---------------------------------------------------------------------------

@dataclass
class RoundGuardrails:
    """Tracks tool usage and failures within one agent round.

    Provides:
    - Per-tool-type call budget enforcement
    - Same-error retry detection and limiting
    - Failure history for structured recovery hints
    """

    # Budget: max calls per tool name in one round
    tool_budget: dict[str, int] = field(default_factory=lambda: {
        "exec_command": 15,
        "write_file": 10,
        "read_file": 20,
        "memory_search": 5,
        "memory_get": 10,
        "web_search": 5,
        "web_fetch": 5,
        "message": 4,
        "send_file": 4,
        "spawn": 3,
        "cron": 5,
    })
    # Max consecutive same-error retries before forced stop
    max_same_error_retries: int = 3
    # Global max tool calls per round
    max_total_tool_calls: int = 40

    # --- Internal tracking ---
    _call_counts: Counter = field(default_factory=Counter)
    _failures: list[ToolFailure] = field(default_factory=list)
    _consecutive_same_errors: int = field(default=0)
    _last_error_sig: str = field(default="")
    _total_calls: int = field(default=0)
    _budget_warnings: set = field(default_factory=set)

    def reset(self) -> None:
        """Reset for a new round."""
        self._call_counts.clear()
        self._failures.clear()
        self._consecutive_same_errors = 0
        self._last_error_sig = ""
        self._total_calls = 0
        self._budget_warnings.clear()

    # --- Check on tool start ---

    def check_tool_allowed(self, tool_name: str, tool_args: dict[str, Any] | None = None) -> str | None:
        """Check if a tool call should be allowed.

        Returns None if allowed, or a reason string if blocked.
        """
        self._total_calls += 1
        self._call_counts[tool_name] += 1

        # Global budget
        if self._total_calls > self.max_total_tool_calls:
            msg = (
                f"Round tool call limit reached ({self.max_total_tool_calls}). "
                "Stop and summarize what was accomplished so far."
            )
            logger.warning("Guardrail: {}", msg)
            return msg

        # Per-tool budget
        budget = self.tool_budget.get(tool_name)
        if budget is not None and self._call_counts[tool_name] > budget:
            msg = (
                f"Tool '{tool_name}' exceeded its per-round budget ({budget} calls). "
                f"Try a different approach or ask the user for guidance."
            )
            if tool_name not in self._budget_warnings:
                logger.warning("Guardrail: {}", msg)
                self._budget_warnings.add(tool_name)
            return msg

        # Same-error retry check
        if self._consecutive_same_errors >= self.max_same_error_retries:
            msg = (
                f"Same error repeated {self._consecutive_same_errors} times. "
                "Stop retrying and try a fundamentally different approach, "
                "or report the issue to the user."
            )
            logger.warning("Guardrail: {}", msg)
            return msg

        return None

    # --- Record tool result ---

    def record_tool_result(
        self,
        tool_name: str,
        tool_args: dict[str, Any] | None,
        result_text: str | None,
        is_error: bool,
    ) -> None:
        """Record a tool call outcome for failure tracking."""
        if not is_error or not result_text:
            # Successful call resets consecutive error counter
            if not is_error:
                self._consecutive_same_errors = 0
                self._last_error_sig = ""
            return

        failure_type = FailureType.classify(result_text)
        args_sig = self._args_signature(tool_name, tool_args)
        error_sig = f"{tool_name}:{failure_type.value}:{args_sig}"

        self._failures.append(ToolFailure(
            tool_name=tool_name,
            failure_type=failure_type,
            error_text=result_text[:200],
            args_signature=args_sig,
        ))

        if error_sig == self._last_error_sig:
            self._consecutive_same_errors += 1
        else:
            self._consecutive_same_errors = 1
            self._last_error_sig = error_sig

        logger.debug(
            "Tool failure: {} type={} consecutive={}/{}",
            tool_name, failure_type.value,
            self._consecutive_same_errors, self.max_same_error_retries,
        )

    # --- Recovery hints ---

    def get_recovery_hint(self) -> str | None:
        """Generate a recovery hint based on recent failures, if any."""
        if not self._failures:
            return None

        recent = self._failures[-3:]
        types = {f.failure_type for f in recent}

        if FailureType.TIMEOUT in types:
            return "[Guardrail] Recent timeouts detected. Consider: shorter operations, chunked processing, or checking service availability."
        if FailureType.PERMISSION in types:
            return "[Guardrail] Permission errors detected. Verify path/access rights before retrying."
        if FailureType.EMPTY_RESULT in types:
            return "[Guardrail] Empty results detected. Try broader search terms or verify the target exists."
        if FailureType.CONFLICT in types:
            return "[Guardrail] Resource conflicts detected. Check for locks or concurrent operations."
        if len(recent) >= 3 and len({f.tool_name for f in recent}) == 1:
            return f"[Guardrail] Tool '{recent[0].tool_name}' has failed {len(recent)} times. Switch to a different approach."

        return None

    # --- Stats ---

    @property
    def stats(self) -> dict[str, Any]:
        return {
            "total_calls": self._total_calls,
            "call_counts": dict(self._call_counts),
            "failure_count": len(self._failures),
            "consecutive_same_errors": self._consecutive_same_errors,
        }

    @staticmethod
    def _args_signature(tool_name: str, args: dict[str, Any] | None) -> str:
        """Create a simplified fingerprint of tool args for dedup."""
        if not args:
            return ""
        # Use first arg value as signature (usually the most distinctive)
        first_val = next(iter(args.values()), "")
        if isinstance(first_val, str):
            return first_val[:60]
        return str(first_val)[:60]
