# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Structured Self-Correction for the Agent Tool Loop

When a tool call fails, instead of passing the raw error string back
to the LLM (which often leads to repeated failures), this module:

    1. Classifies the error into actionable categories
    2. Tracks per-tool retry counts within a single conversation turn
    3. Formats error feedback with reflection prompts that guide
       the LLM toward recovery

Usage:
    tracker = ToolRetryTracker()

    result = execute_tool(...)
    if tracker.is_error(result):
        error_class = tracker.classify_error(result)
        if tracker.should_retry(tool_name):
            result = tracker.format_error_feedback(tool_name, result, error_class)
        else:
            result = tracker.format_give_up(tool_name, result, error_class)
"""

import logging
import re
from dataclasses import dataclass
from enum import Enum
from typing import Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)


class ErrorClass(str, Enum):
    """Categories of tool errors with different recovery strategies."""

    PARAMETER_ERROR = "parameter_error"
    PERMISSION_DENIED = "permission_denied"
    NOT_FOUND = "not_found"
    SERVICE_DOWN = "service_down"
    VALIDATION_ERROR = "validation_error"
    UNKNOWN_TOOL = "unknown_tool"
    RATE_LIMIT = "rate_limit"
    UNKNOWN = "unknown"


_ERROR_PATTERNS: List[Tuple[re.Pattern, ErrorClass]] = [
    (
        re.compile(r"missing required|required field|required parameter|missing.*argument", re.I),
        ErrorClass.PARAMETER_ERROR,
    ),
    (
        re.compile(r"invalid.*type|expected.*got|type error|must be.*not", re.I),
        ErrorClass.PARAMETER_ERROR,
    ),
    (
        re.compile(r"invalid.*value|out of range|must be positive|cannot be negative", re.I),
        ErrorClass.VALIDATION_ERROR,
    ),
    (
        re.compile(r"invalid.*format|malformed|bad.*format|parse error", re.I),
        ErrorClass.VALIDATION_ERROR,
    ),
    (
        re.compile(r"permission denied|not authorized|forbidden|access denied", re.I),
        ErrorClass.PERMISSION_DENIED,
    ),
    (
        re.compile(r"⚠️.*permission|missing.*capability|read.?only", re.I),
        ErrorClass.PERMISSION_DENIED,
    ),
    (
        re.compile(r"capability.*not granted|build trust to unlock", re.I),
        ErrorClass.PERMISSION_DENIED,
    ),
    (
        re.compile(r"not found|no.*match|does not exist|no results|no.*found", re.I),
        ErrorClass.NOT_FOUND,
    ),
    (re.compile(r"unknown tool|no tool named", re.I), ErrorClass.UNKNOWN_TOOL),
    (
        re.compile(r"connection refused|connection error|timed? ?out|unreachable", re.I),
        ErrorClass.SERVICE_DOWN,
    ),
    (re.compile(r"service unavailable|503|502|500.*error", re.I), ErrorClass.SERVICE_DOWN),
    (re.compile(r"ollama.*not running|server.*down|api.*error", re.I), ErrorClass.SERVICE_DOWN),
    (re.compile(r"rate limit|too many requests|429|throttl", re.I), ErrorClass.RATE_LIMIT),
]


def classify_error(error_text: str) -> ErrorClass:
    """Classify a tool error into an actionable category."""
    for pattern, error_class in _ERROR_PATTERNS:
        if pattern.search(error_text):
            return error_class
    return ErrorClass.UNKNOWN


_RECOVERY_GUIDANCE = {
    ErrorClass.PARAMETER_ERROR: (
        "This failed because of a parameter issue. "
        "Check which parameter was wrong and retry with corrected values."
    ),
    ErrorClass.VALIDATION_ERROR: (
        "The data format or value was invalid. "
        "Review the expected format and retry with valid data."
    ),
    ErrorClass.PERMISSION_DENIED: (
        "This action is not permitted for the current user. "
        "Inform the user about the permission limitation. Do not retry."
    ),
    ErrorClass.NOT_FOUND: (
        "The requested resource was not found. "
        "Try searching with different terms, check for typos, "
        "or ask the user to clarify what they're looking for."
    ),
    ErrorClass.UNKNOWN_TOOL: (
        "This tool does not exist. Check available tools and use one that matches your intent."
    ),
    ErrorClass.SERVICE_DOWN: (
        "An external service is unavailable. "
        "Inform the user and suggest trying again later. Do not retry."
    ),
    ErrorClass.RATE_LIMIT: ("Rate limit hit. Inform the user and suggest waiting before retrying."),
    ErrorClass.UNKNOWN: (
        "Analyze the error message to understand what went wrong. "
        "Decide whether to retry with a different approach or inform the user."
    ),
}

_RETRYABLE = {
    ErrorClass.PARAMETER_ERROR,
    ErrorClass.VALIDATION_ERROR,
    ErrorClass.NOT_FOUND,
    ErrorClass.UNKNOWN_TOOL,
    ErrorClass.UNKNOWN,
}


def is_retryable(error_class: ErrorClass) -> bool:
    """Whether this error class is worth retrying."""
    return error_class in _RETRYABLE


@dataclass
class ToolFailure:
    """Record of a single tool failure."""

    tool_name: str
    error_text: str
    error_class: ErrorClass
    attempt: int


class ToolRetryTracker:
    """
    Tracks tool failures within a single conversation turn.
    Create a new instance for each chat() call.
    """

    DEFAULT_MAX_RETRIES = 2

    def __init__(self, max_retries: int = DEFAULT_MAX_RETRIES):
        self.max_retries = max_retries
        self._attempts: Dict[str, int] = {}
        self._failures: List[ToolFailure] = []
        self._error_classes: Dict[str, ErrorClass] = {}

    def is_error(self, result: str) -> bool:
        """Check if a tool result is an error."""
        if not result:
            return False
        return (
            result.startswith("Error") or result.startswith("⚠️") or "error" in result[:50].lower()
        )

    def record_failure(self, tool_name: str, error_text: str) -> ErrorClass:
        """Record a tool failure and classify it."""
        error_class = classify_error(error_text)
        attempt = self._attempts.get(tool_name, 0) + 1
        self._attempts[tool_name] = attempt
        self._error_classes[tool_name] = error_class
        self._failures.append(
            ToolFailure(
                tool_name=tool_name,
                error_text=error_text,
                error_class=error_class,
                attempt=attempt,
            )
        )
        logger.info(
            f"Tool failure: {tool_name} attempt={attempt}/{self.max_retries} "
            f"class={error_class.value}"
        )
        return error_class

    def should_retry(self, tool_name: str) -> bool:
        """Whether this tool should be retried."""
        attempts = self._attempts.get(tool_name, 0)
        if attempts >= self.max_retries:
            return False
        error_class = self._error_classes.get(tool_name, ErrorClass.UNKNOWN)
        return is_retryable(error_class)

    def get_attempts(self, tool_name: str) -> int:
        """Get the number of failed attempts for a tool."""
        return self._attempts.get(tool_name, 0)

    def format_error_feedback(
        self,
        tool_name: str,
        error_text: str,
        error_class: Optional[ErrorClass] = None,
    ) -> str:
        """Format a tool error with structured reflection guidance."""
        if error_class is None:
            error_class = self._error_classes.get(tool_name, ErrorClass.UNKNOWN)
        attempts = self._attempts.get(tool_name, 1)
        guidance = _RECOVERY_GUIDANCE.get(error_class, _RECOVERY_GUIDANCE[ErrorClass.UNKNOWN])
        return (
            f"TOOL CALL FAILED: {tool_name}\n"
            f"Error: {error_text}\n"
            f"Error type: {error_class.value}\n"
            f"Attempt: {attempts}/{self.max_retries}\n"
            f"\n"
            f"Recovery: {guidance}"
        )

    def format_give_up(
        self,
        tool_name: str,
        error_text: str,
        error_class: Optional[ErrorClass] = None,
    ) -> str:
        """Format the message when max retries are exhausted."""
        if error_class is None:
            error_class = self._error_classes.get(tool_name, ErrorClass.UNKNOWN)
        attempts = self._attempts.get(tool_name, 0)
        return (
            f"TOOL CALL FAILED (max retries reached): {tool_name}\n"
            f"Error: {error_text}\n"
            f"Error type: {error_class.value}\n"
            f"Attempts: {attempts}/{self.max_retries}\n"
            f"\n"
            f"Do NOT retry this tool. Inform the user about the issue "
            f"and suggest what they can do to resolve it."
        )

    def get_retry_delay(self, tool_name: str) -> float:
        """Return delay (seconds) before retrying a failed tool.

        SERVICE_DOWN / RATE_LIMIT get exponential backoff (2s base, max 30s).
        All other errors return 0.0 — the LLM will adjust params, no sleep needed.
        """
        error_class = self._error_classes.get(tool_name)
        if error_class not in (ErrorClass.SERVICE_DOWN, ErrorClass.RATE_LIMIT):
            return 0.0
        attempts = self._attempts.get(tool_name, 1)
        return min(2.0 * (2 ** (attempts - 1)), 30.0)

    def get_summary(self) -> Dict:
        """Summary of all failures this turn."""
        return {
            "total_failures": len(self._failures),
            "tools_failed": list(self._attempts.keys()),
            "retries_exhausted": [
                name for name, count in self._attempts.items() if count >= self.max_retries
            ],
            "error_classes": {name: cls.value for name, cls in self._error_classes.items()},
        }
