"""VERONICA Retry Containment - Budget-aware retry wrapper."""

from __future__ import annotations

__all__ = ["RetryContainer"]
from dataclasses import dataclass, field
from typing import Callable, TypeVar, Optional, Any
import random
import threading
import time
import logging

from veronica_core.runtime_policy import PolicyContext, PolicyDecision

logger = logging.getLogger(__name__)

T = TypeVar("T")


@dataclass
class RetryContainer:
    """Budget-aware retry wrapper for LLM calls.

    Unlike per-call retry limits, this enforces a total retry budget
    across the entire request chain. Prevents 3 retries x 5 nested
    calls = 15 LLM calls from one user action.

    C1 CONCURRENCY WARNING: ``RetryContainer`` is NOT safe for concurrent
    ``execute()`` calls on the same instance. ``_attempt_count`` and
    ``_total_retries`` are instance-level state; concurrent callers will race
    on these counters and produce incorrect values. Each call chain should use
    its own ``RetryContainer`` instance. This is by design -- the container
    tracks per-chain retry budget, not per-thread.

    Example:
        retry = RetryContainer(max_retries=3, backoff_base=1.0)
        result = retry.execute(call_llm_api, prompt="hello")
    """

    max_retries: int = 3
    backoff_base: float = 1.0
    backoff_max: float = 30.0
    # jitter fraction applied to each backoff delay (0.0 = no jitter, 1.0 = ±100%).
    # Default 0.25 adds ±25% randomness to prevent thundering herd when many
    # RetryContainer instances retry simultaneously after a shared failure.
    jitter: float = 0.25

    _attempt_count: int = field(default=0, init=False)
    _total_retries: int = field(default=0, init=False)
    _last_error: Optional[Exception] = field(default=None, init=False, repr=False)
    _lock: threading.Lock = field(
        default_factory=threading.Lock, init=False, repr=False
    )

    def execute(
        self,
        fn: Callable[..., T],
        *args: Any,
        **kwargs: Any,
    ) -> T:
        """Execute function with retry containment.

        Args:
            fn: Callable to execute
            *args: Positional arguments
            **kwargs: Keyword arguments

        Returns:
            Result from fn

        Raises:
            Exception: Last exception if all retries exhausted
        """
        with self._lock:
            self._attempt_count = 0

        for attempt in range(self.max_retries + 1):
            with self._lock:
                self._attempt_count = attempt + 1

            try:
                result = fn(*args, **kwargs)
                with self._lock:
                    self._last_error = None  # Clear error state on success
                if attempt > 0:
                    logger.info(f"[VERONICA_RETRY] Succeeded on attempt {attempt + 1}")
                return result

            except Exception as e:
                with self._lock:
                    self._last_error = e
                    self._total_retries += 1

                if attempt >= self.max_retries:
                    logger.warning(
                        f"[VERONICA_RETRY] All {self.max_retries} retries exhausted: {type(e).__name__}"
                    )
                    raise

                base_delay = self.backoff_base * (2**attempt)
                if self.jitter > 0.0:
                    base_delay *= 1.0 + random.uniform(-self.jitter, self.jitter)
                delay = min(max(0.0, base_delay), self.backoff_max)
                logger.info(
                    f"[VERONICA_RETRY] Attempt {attempt + 1} failed: {type(e).__name__}. "
                    f"Retrying in {delay:.1f}s "
                    f"({self.max_retries - attempt} remaining)"
                )
                # Release lock before sleeping so other threads can call
                # reset() / check() without being blocked during backoff.
                time.sleep(delay)

        # Unreachable under normal operation (max_retries >= 0).
        # The loop always exits via return (success) or raise (last attempt).
        raise RuntimeError("max retries exceeded")  # pragma: no cover

    @property
    def attempt_count(self) -> int:
        """Number of attempts in the last execute() call."""
        with self._lock:
            return self._attempt_count

    @property
    def total_retries(self) -> int:
        """Total retries across all execute() calls."""
        with self._lock:
            return self._total_retries

    @property
    def last_error(self) -> Optional[Exception]:
        """Last error encountered."""
        with self._lock:
            return self._last_error

    @property
    def policy_type(self) -> str:
        """RuntimePolicy protocol: policy type identifier."""
        return "retry_budget"

    def check(self, context: PolicyContext) -> PolicyDecision:
        """RuntimePolicy protocol: check if retry budget is available.

        Returns denied if the container is in an error state
        (last execution failed after exhausting all retries).
        Does NOT modify retry state.

        Args:
            context: PolicyContext (fields unused)

        Returns:
            PolicyDecision allowing or denying the operation
        """
        with self._lock:
            last_error = self._last_error
            total_retries = self._total_retries
        if last_error is not None:
            return PolicyDecision(
                allowed=False,
                policy_type=self.policy_type,
                reason=(f"Retry budget exhausted ({total_retries} retries used)"),
            )
        return PolicyDecision(allowed=True, policy_type=self.policy_type)

    def record_failure(self, error: Optional[Exception] = None) -> None:
        """Record a single failure without calling execute().

        Called by framework adapter error hooks (on_llm_error, LLMCallFailedEvent)
        to increment the retry counter and set last_error so that subsequent
        check() calls reflect the failure state. Use reset() to clear.

        WARNING: Do NOT call record_failure() from inside a function passed to
        execute(). execute() already tracks failures internally. Calling both
        will double-count failures and exhaust the retry budget prematurely.

        Args:
            error: The exception to record, or None for a generic failure.
        """
        sentinel = (
            error if error is not None else RuntimeError("adapter-recorded failure")
        )
        with self._lock:
            self._total_retries += 1
            total = self._total_retries
            if self._total_retries > self.max_retries:
                self._last_error = sentinel
        logger.debug(
            "[VERONICA_RETRY] record_failure(): total_retries=%d", total
        )

    def reset(self) -> None:
        """Reset retry state for reuse."""
        with self._lock:
            self._attempt_count = 0
            self._total_retries = 0
            self._last_error = None
        logger.info("[VERONICA_RETRY] Retry state reset")
