"""
Zubbl Failure Interceptor — Real-Time Failure Recovery

This is the core wedge product: when an agent fails, Zubbl catches it
and applies a fix automatically, without code changes or redeployment.

The flow:
    1. Agent runs → fails (exception, timeout, wrong answer)
    2. Interceptor catches the failure IN REAL-TIME
    3. Matches against known failure patterns
    4. Applies the best fix: retry with modified prompt, fallback strategy, or cached fix
    5. If fix works → learns and pre-applies for future runs
    6. If no fix available → records the failure pattern for learning

This turns Zubbl from "learn over time" into "fix it right now."

Example:
    from zubbl import ZubblClient

    zubbl = ZubblClient(api_key="your-key")
    agent = zubbl.wrap(your_agent, auto_recover=True)

    # Agent fails internally → Zubbl retries with a fix → returns correct result
    # No code change. No redeploy. The customer never sees the failure.
"""

import hashlib
import time
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Callable, Dict, List, Optional, Tuple


class FixType(str, Enum):
    """Types of automatic fixes Zubbl can apply"""

    RETRY = "retry"  # Simple retry (transient errors)
    RETRY_WITH_BACKOFF = "retry_backoff"  # Retry with exponential backoff
    PROMPT_PATCH = "prompt_patch"  # Modify the prompt/input
    FALLBACK_MODEL = "fallback_model"  # Try a different model/provider
    CACHED_RESPONSE = "cached_response"  # Return a known-good cached response
    INPUT_TRANSFORM = "input_transform"  # Transform input to avoid failure
    DECOMPOSE = "decompose"  # Break task into smaller subtasks
    SKIP = "skip"  # Skip this step, continue pipeline


class FailureCategory(str, Enum):
    """Categories of agent failures"""

    EXCEPTION = "exception"  # Python exception
    TIMEOUT = "timeout"  # Exceeded time limit
    RATE_LIMIT = "rate_limit"  # API rate limit hit
    CONTEXT_OVERFLOW = "context_overflow"  # Token limit exceeded
    HALLUCINATION = "hallucination"  # Confident but wrong (detected post-hoc)
    EMPTY_RESPONSE = "empty_response"  # Agent returned nothing
    MALFORMED_OUTPUT = "malformed_output"  # Output doesn't match expected format
    API_ERROR = "api_error"  # External API returned error
    PARTIAL_FAILURE = "partial_failure"  # Some steps succeeded, some failed


@dataclass
class FailurePattern:
    """A recognized pattern of agent failure"""

    pattern_id: str
    category: FailureCategory
    error_signature: str  # Normalized error fingerprint
    task_type: str  # What kind of task fails this way
    occurrence_count: int = 0
    last_seen: Optional[datetime] = None
    fixes_tried: Dict[str, int] = field(default_factory=dict)  # fix_type → attempt count
    fixes_succeeded: Dict[str, int] = field(default_factory=dict)  # fix_type → success count
    best_fix: Optional[str] = None
    best_fix_success_rate: float = 0.0
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class Fix:
    """A fix to apply when a failure is detected"""

    fix_type: FixType
    confidence: float = 0.0  # How likely this fix will work (0-1)
    prompt_patch: Optional[str] = None  # Modified prompt/instruction
    fallback_model: Optional[str] = None  # Alternative model to try
    cached_response: Optional[Any] = None  # Known-good response
    input_transform: Optional[Callable] = None  # Function to transform input
    max_retries: int = 3
    backoff_base: float = 1.0  # Base seconds for exponential backoff
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class RecoveryResult:
    """Result of an auto-recovery attempt"""

    success: bool
    fix_applied: Optional[FixType] = None
    attempts: int = 0
    original_error: Optional[str] = None
    recovered_output: Optional[Any] = None
    latency_ms: int = 0
    pattern_id: Optional[str] = None


class FailureInterceptor:
    """
    Real-time failure interception and auto-recovery.

    This is the engine that makes "Your agent fails, Zubbl fixes it" real.

    It maintains an in-memory cache of failure patterns + fixes,
    synced with the server periodically. When a failure happens:

    1. Classify the failure (exception type, error message, task context)
    2. Match against known patterns
    3. Apply the best fix from pattern history
    4. If no known fix, try generic strategies (retry, backoff, fallback)
    5. Record result for learning
    """

    # Error signatures that indicate rate limiting
    RATE_LIMIT_ERRORS = [
        "rate limit",
        "429",
        "too many requests",
        "quota exceeded",
    ]

    # Error signatures that indicate transient API/server errors (safe to retry)
    TRANSIENT_ERRORS = [
        "503",
        "502",
        "500",
        "connection reset",
        "temporarily unavailable",
        "server error",
        "internal server error",
        "overloaded",
        "capacity",
        "try again",
        "retry",
        "connection refused",
    ]

    # Error signatures that indicate context/token overflow
    CONTEXT_ERRORS = [
        "context length",
        "maximum context",
        "token limit",
        "too many tokens",
        "context_length_exceeded",
        "max_tokens",
        "input too long",
    ]

    def __init__(
        self,
        max_retries: int = 3,
        retry_backoff: float = 1.0,
        enable_prompt_patching: bool = True,
        enable_fallback: bool = True,
        fallback_models: Optional[List[str]] = None,
    ):
        self.max_retries = max_retries
        self.retry_backoff = retry_backoff
        self.enable_prompt_patching = enable_prompt_patching
        self.enable_fallback = enable_fallback
        self.fallback_models = fallback_models or []

        # In-memory pattern cache (synced from server)
        self._patterns: Dict[str, FailurePattern] = {}
        # Response cache: task_hash → known-good response
        self._response_cache: Dict[str, Any] = {}
        # Prompt patches: task_type → list of patches that worked
        self._prompt_patches: Dict[str, List[Dict]] = {}

    def intercept_and_recover(
        self,
        func: Callable,
        args: tuple,
        kwargs: dict,
        task_description: str = "",
        task_type: str = "",
    ) -> RecoveryResult:
        """
        Execute a function with automatic failure recovery.

        This is the main entry point called by ZubblWrappedAgent.
        If the function fails, it tries to recover automatically.

        Args:
            func: The agent method to call
            args: Positional arguments
            kwargs: Keyword arguments
            task_description: Human-readable task description
            task_type: Normalized task type for pattern matching

        Returns:
            RecoveryResult with success status and output
        """
        start_time = time.time()
        original_error = None
        attempts = 0

        # Step 1: Try the original call
        try:
            result = func(*args, **kwargs)

            # Check for empty/None responses
            if result is None or (isinstance(result, str) and not result.strip()):
                original_error = "empty_response"
                category = FailureCategory.EMPTY_RESPONSE
            else:
                # Success on first try
                self._cache_good_response(task_description, result)
                return RecoveryResult(
                    success=True,
                    attempts=1,
                    recovered_output=result,
                    latency_ms=int((time.time() - start_time) * 1000),
                )
        except Exception as e:
            original_error = str(e)
            category = self._classify_error(e)
            attempts = 1

        # Step 2: Failure detected — find the best fix
        error_sig = self._compute_error_signature(original_error, category)
        pattern = self._find_matching_pattern(error_sig, task_type)

        # Step 3: Build fix strategy (ordered by confidence)
        fixes = self._build_fix_strategy(
            category=category,
            pattern=pattern,
            task_description=task_description,
            task_type=task_type,
            original_error=original_error,
        )

        # Step 4: Try fixes in order
        for fix in fixes:
            attempts += 1
            try:
                result = self._apply_fix(fix, func, args, kwargs, task_description)

                if result is not None:
                    # Fix worked!
                    elapsed = int((time.time() - start_time) * 1000)
                    self._record_fix_success(error_sig, task_type, fix, category)
                    self._cache_good_response(task_description, result)

                    return RecoveryResult(
                        success=True,
                        fix_applied=fix.fix_type,
                        attempts=attempts,
                        original_error=original_error,
                        recovered_output=result,
                        latency_ms=elapsed,
                        pattern_id=pattern.pattern_id if pattern else None,
                    )
            except Exception:
                # This fix didn't work, try the next one
                continue

        # Step 5: All fixes failed — record the failure for learning
        elapsed = int((time.time() - start_time) * 1000)
        self._record_failure(error_sig, task_type, category, original_error)

        return RecoveryResult(
            success=False,
            attempts=attempts,
            original_error=original_error,
            latency_ms=elapsed,
            pattern_id=pattern.pattern_id if pattern else None,
        )

    # =========================================================================
    # Error Classification
    # =========================================================================

    def _classify_error(self, error: Exception) -> FailureCategory:
        """Classify an exception into a failure category"""
        error_str = str(error).lower()
        error_type = type(error).__name__.lower()

        # Rate limits (check first — most specific)
        if any(sig in error_str for sig in self.RATE_LIMIT_ERRORS):
            return FailureCategory.RATE_LIMIT

        # Context overflow
        if any(sig in error_str for sig in self.CONTEXT_ERRORS):
            return FailureCategory.CONTEXT_OVERFLOW

        # Timeouts
        if "timeout" in error_str or "timed out" in error_str or "timeout" in error_type:
            return FailureCategory.TIMEOUT

        # API/server errors (transient)
        if any(sig in error_str for sig in self.TRANSIENT_ERRORS):
            return FailureCategory.API_ERROR

        # Generic exception
        return FailureCategory.EXCEPTION

    def _compute_error_signature(self, error_msg: str, category: FailureCategory) -> str:
        """
        Compute a normalized error fingerprint.

        This is critical for pattern matching — we need to recognize
        "the same error in a slightly different form."

        Strategy:
        - Strip variable parts (IDs, timestamps, file paths, line numbers)
        - Keep the structural error message
        - Hash for stable lookup
        """
        import re

        normalized = error_msg.lower().strip()

        # Remove variable parts (order matters — specific patterns first)
        normalized = re.sub(r"0x[0-9a-f]+", "<addr>", normalized)  # Memory addresses
        normalized = re.sub(
            r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}", "<uuid>", normalized
        )  # UUIDs
        normalized = re.sub(
            r"\d{4}-\d{2}-\d{2}[T ]\d{2}:\d{2}:\d{2}", "<timestamp>", normalized
        )  # Timestamps
        normalized = re.sub(r"\b\d+\.\d+\.\d+\.\d+\b", "<ip>", normalized)  # IPs
        normalized = re.sub(r"/[\w/.-]+\.\w+", "<filepath>", normalized)  # File paths
        normalized = re.sub(r"request id:?\s*\S+", "request id: <id>", normalized)  # Request IDs
        normalized = re.sub(r"line \d+", "line <N>", normalized)  # Line numbers
        normalized = re.sub(r"\b\d+ms\b", "<time>", normalized)  # Latencies
        normalized = re.sub(r"\b\d+\b", "<N>", normalized)  # All remaining numbers

        # Create signature from category + normalized message
        sig_input = f"{category.value}:{normalized}"
        return hashlib.md5(sig_input.encode()).hexdigest()[:16]

    # =========================================================================
    # Pattern Matching
    # =========================================================================

    def _find_matching_pattern(self, error_sig: str, task_type: str) -> Optional[FailurePattern]:
        """Find a known failure pattern that matches this error"""
        # Exact signature match (most reliable)
        if error_sig in self._patterns:
            return self._patterns[error_sig]

        # Task-type + category match (broader)
        for pid, pattern in self._patterns.items():
            if pattern.task_type == task_type and pattern.occurrence_count > 2:
                return pattern

        return None

    # =========================================================================
    # Fix Strategy
    # =========================================================================

    def _build_fix_strategy(
        self,
        category: FailureCategory,
        pattern: Optional[FailurePattern],
        task_description: str,
        task_type: str,
        original_error: str,
    ) -> List[Fix]:
        """
        Build an ordered list of fixes to try.

        If we have a known pattern with a proven fix, use that first.
        Otherwise, fall back to category-based generic strategies.
        """
        fixes = []

        # 1. Known fix from pattern history (highest confidence)
        if pattern and pattern.best_fix:
            fix_type = FixType(pattern.best_fix)
            fixes.append(
                Fix(
                    fix_type=fix_type,
                    confidence=pattern.best_fix_success_rate,
                    max_retries=1,
                    metadata={"source": "pattern_history", "pattern_id": pattern.pattern_id},
                )
            )

        # 2. Cached response (if available for this task)
        task_hash = self._hash_task(task_description)
        if task_hash in self._response_cache:
            fixes.append(
                Fix(
                    fix_type=FixType.CACHED_RESPONSE,
                    confidence=0.8,
                    cached_response=self._response_cache[task_hash],
                    metadata={"source": "response_cache"},
                )
            )

        # 3. Category-specific fixes
        category_fixes = self._get_category_fixes(category, task_description, original_error)
        fixes.extend(category_fixes)

        # 4. Prompt patching (if enabled and we have patches)
        if self.enable_prompt_patching and task_type in self._prompt_patches:
            for patch in self._prompt_patches[task_type]:
                fixes.append(
                    Fix(
                        fix_type=FixType.PROMPT_PATCH,
                        confidence=patch.get("success_rate", 0.3),
                        prompt_patch=patch.get("patch"),
                        metadata={"source": "learned_patch"},
                    )
                )

        # Sort by confidence (highest first)
        fixes.sort(key=lambda f: f.confidence, reverse=True)

        # Limit to max_retries
        return fixes[: self.max_retries]

    def _get_category_fixes(
        self,
        category: FailureCategory,
        task_description: str,
        original_error: str,
    ) -> List[Fix]:
        """Get generic fixes based on failure category"""
        fixes = []

        if category == FailureCategory.RATE_LIMIT:
            # Rate limits → exponential backoff retry
            fixes.append(
                Fix(
                    fix_type=FixType.RETRY_WITH_BACKOFF,
                    confidence=0.9,
                    max_retries=3,
                    backoff_base=2.0,
                    metadata={"reason": "rate_limit_detected"},
                )
            )

        elif category == FailureCategory.TIMEOUT:
            # Timeout → retry once, then try decomposition
            fixes.append(
                Fix(
                    fix_type=FixType.RETRY,
                    confidence=0.4,
                    max_retries=1,
                    metadata={"reason": "timeout_retry"},
                )
            )
            fixes.append(
                Fix(
                    fix_type=FixType.DECOMPOSE,
                    confidence=0.3,
                    metadata={"reason": "task_too_complex"},
                )
            )

        elif category == FailureCategory.CONTEXT_OVERFLOW:
            # Context overflow → truncate input and retry
            fixes.append(
                Fix(
                    fix_type=FixType.INPUT_TRANSFORM,
                    confidence=0.7,
                    input_transform=self._truncate_input,
                    metadata={"reason": "context_overflow"},
                )
            )

        elif category == FailureCategory.EMPTY_RESPONSE:
            # Empty response → retry with prompt reinforcement
            fixes.append(
                Fix(
                    fix_type=FixType.PROMPT_PATCH,
                    confidence=0.5,
                    prompt_patch="Please provide a complete, non-empty response.",
                    metadata={"reason": "empty_response_reinforcement"},
                )
            )
            fixes.append(
                Fix(
                    fix_type=FixType.RETRY,
                    confidence=0.4,
                    max_retries=2,
                    metadata={"reason": "empty_response_retry"},
                )
            )

        elif category in (FailureCategory.API_ERROR, FailureCategory.EXCEPTION):
            # Generic errors → retry with backoff, then fallback model
            fixes.append(
                Fix(
                    fix_type=FixType.RETRY_WITH_BACKOFF,
                    confidence=0.5,
                    max_retries=2,
                    backoff_base=1.0,
                    metadata={"reason": "generic_retry"},
                )
            )
            if self.enable_fallback and self.fallback_models:
                for model in self.fallback_models:
                    fixes.append(
                        Fix(
                            fix_type=FixType.FALLBACK_MODEL,
                            confidence=0.4,
                            fallback_model=model,
                            metadata={"reason": "fallback_to_" + model},
                        )
                    )

        return fixes

    # =========================================================================
    # Fix Application
    # =========================================================================

    def _apply_fix(
        self,
        fix: Fix,
        func: Callable,
        args: tuple,
        kwargs: dict,
        task_description: str,
    ) -> Any:
        """Apply a specific fix and return the result"""

        if fix.fix_type == FixType.RETRY:
            return func(*args, **kwargs)

        elif fix.fix_type == FixType.RETRY_WITH_BACKOFF:
            for attempt in range(fix.max_retries):
                try:
                    wait_time = fix.backoff_base * (2**attempt)
                    time.sleep(wait_time)
                    return func(*args, **kwargs)
                except Exception:
                    if attempt == fix.max_retries - 1:
                        raise
            return None

        elif fix.fix_type == FixType.CACHED_RESPONSE:
            return fix.cached_response

        elif fix.fix_type == FixType.PROMPT_PATCH:
            # Modify the first string argument (likely the prompt/task)
            patched_args = list(args)
            if patched_args and isinstance(patched_args[0], str):
                patched_args[0] = f"{patched_args[0]}\n\n{fix.prompt_patch}"
            elif "prompt" in kwargs:
                kwargs["prompt"] = f"{kwargs['prompt']}\n\n{fix.prompt_patch}"
            elif "task" in kwargs:
                kwargs["task"] = f"{kwargs['task']}\n\n{fix.prompt_patch}"
            elif "messages" in kwargs and isinstance(kwargs["messages"], list):
                # Append to the last user message
                for msg in reversed(kwargs["messages"]):
                    if isinstance(msg, dict) and msg.get("role") == "user":
                        msg["content"] = f"{msg['content']}\n\n{fix.prompt_patch}"
                        break
            return func(*tuple(patched_args), **kwargs)

        elif fix.fix_type == FixType.INPUT_TRANSFORM:
            if fix.input_transform:
                new_args, new_kwargs = fix.input_transform(args, kwargs)
                return func(*new_args, **new_kwargs)
            return func(*args, **kwargs)

        elif fix.fix_type == FixType.FALLBACK_MODEL:
            # Try with a different model — inject model name into kwargs
            kwargs_copy = dict(kwargs)
            kwargs_copy["model"] = fix.fallback_model
            try:
                return func(*args, **kwargs_copy)
            except TypeError:
                # Function doesn't accept model kwarg — try original
                return func(*args, **kwargs)

        elif fix.fix_type == FixType.DECOMPOSE:
            # Decomposition is a no-op at Phase 1 — would need LLM call
            # to break task into subtasks. For now, just retry.
            return func(*args, **kwargs)

        elif fix.fix_type == FixType.SKIP:
            return None

        return func(*args, **kwargs)

    # =========================================================================
    # Input Transforms
    # =========================================================================

    @staticmethod
    def _truncate_input(args: tuple, kwargs: dict) -> Tuple[tuple, dict]:
        """Truncate input to fit within context limits"""
        new_args = list(args)
        for i, arg in enumerate(new_args):
            if isinstance(arg, str) and len(arg) > 10000:
                # Keep first and last parts, drop the middle
                keep = 4000
                new_args[i] = (
                    arg[:keep]
                    + "\n\n[... content truncated for context limits ...]\n\n"
                    + arg[-keep:]
                )
        return tuple(new_args), kwargs

    # =========================================================================
    # Learning & Caching
    # =========================================================================

    def _cache_good_response(self, task_description: str, response: Any) -> None:
        """Cache a known-good response for fallback use"""
        if not task_description:
            return
        task_hash = self._hash_task(task_description)
        self._response_cache[task_hash] = response

    def _record_fix_success(
        self, error_sig: str, task_type: str, fix: Fix, category: FailureCategory
    ) -> None:
        """Record that a fix worked — updates pattern for future use"""
        if error_sig not in self._patterns:
            self._patterns[error_sig] = FailurePattern(
                pattern_id=f"fp_{error_sig}",
                category=category,
                error_signature=error_sig,
                task_type=task_type,
            )

        pattern = self._patterns[error_sig]
        fix_name = fix.fix_type.value

        pattern.fixes_tried[fix_name] = pattern.fixes_tried.get(fix_name, 0) + 1
        pattern.fixes_succeeded[fix_name] = pattern.fixes_succeeded.get(fix_name, 0) + 1
        pattern.occurrence_count += 1
        pattern.last_seen = datetime.utcnow()

        # Update best fix
        best_rate = 0.0
        for fx, tried in pattern.fixes_tried.items():
            succeeded = pattern.fixes_succeeded.get(fx, 0)
            rate = succeeded / tried if tried > 0 else 0.0
            if rate > best_rate:
                best_rate = rate
                pattern.best_fix = fx
                pattern.best_fix_success_rate = rate

    def _record_failure(
        self, error_sig: str, task_type: str, category: FailureCategory, error_msg: str
    ) -> None:
        """Record an unrecoverable failure — future learning data"""
        if error_sig not in self._patterns:
            self._patterns[error_sig] = FailurePattern(
                pattern_id=f"fp_{error_sig}",
                category=category,
                error_signature=error_sig,
                task_type=task_type,
            )

        pattern = self._patterns[error_sig]
        pattern.occurrence_count += 1
        pattern.last_seen = datetime.utcnow()
        pattern.metadata["last_error"] = error_msg[:500]

    def _hash_task(self, task: str) -> str:
        """Hash a task description for cache lookup"""
        return hashlib.md5(task.lower().strip().encode()).hexdigest()[:12]

    # =========================================================================
    # Pattern Sync (Server <-> Local)
    # =========================================================================

    def load_patterns_from_server(self, patterns_data: List[Dict]) -> None:
        """Load failure patterns from the Zubbl server"""
        for p in patterns_data:
            sig = p.get("error_signature", "")
            self._patterns[sig] = FailurePattern(
                pattern_id=p.get("pattern_id", ""),
                category=FailureCategory(p.get("category", "exception")),
                error_signature=sig,
                task_type=p.get("task_type", ""),
                occurrence_count=p.get("occurrence_count", 0),
                fixes_tried=p.get("fixes_tried", {}),
                fixes_succeeded=p.get("fixes_succeeded", {}),
                best_fix=p.get("best_fix"),
                best_fix_success_rate=p.get("best_fix_success_rate", 0.0),
            )

    def load_prompt_patches(self, patches_data: Dict[str, List[Dict]]) -> None:
        """Load learned prompt patches from the server"""
        self._prompt_patches = patches_data

    def export_patterns(self) -> List[Dict]:
        """Export local patterns for server sync"""
        return [
            {
                "pattern_id": p.pattern_id,
                "category": p.category.value,
                "error_signature": p.error_signature,
                "task_type": p.task_type,
                "occurrence_count": p.occurrence_count,
                "last_seen": p.last_seen.isoformat() if p.last_seen else None,
                "fixes_tried": p.fixes_tried,
                "fixes_succeeded": p.fixes_succeeded,
                "best_fix": p.best_fix,
                "best_fix_success_rate": p.best_fix_success_rate,
                "metadata": p.metadata,
            }
            for p in self._patterns.values()
        ]

    @property
    def pattern_count(self) -> int:
        return len(self._patterns)

    @property
    def cached_responses(self) -> int:
        return len(self._response_cache)
