"""RetryMiddleware — MAF-style composable retry for AgentByte middleware chain."""

import random
from enum import Enum
from typing import Any, List, Optional, Tuple, Type, Union

from .base import BaseMiddleware, MiddlewareContext, OperationType

_DEFAULT_OPERATION_TYPES: frozenset = frozenset(["tool_call", "memory_access"])
_RETRY_MIDDLEWARE_KEY = "_retry_middleware"


class RetryMiddleware(BaseMiddleware):
    """Retry middleware with exponential backoff, inspired by Microsoft Agent Framework.

    Covers ``tool_call`` and ``memory_access`` by default.  ``model_call`` and
    ``embedding_call`` are excluded because their provider clients already implement
    retry with exponential backoff in ``llm/openai.py``, ``llm/azure_openai.py``,
    ``llm/openai_embedding.py``, and ``llm/azure_openai_embedding.py``.  Enabling
    those operations here without setting ``max_retries=0`` on the relevant client
    causes nested retry (e.g. 3 middleware attempts × 3 client retries = 9 calls).

    Streaming tools are out of scope for Phase 1.  The streaming-tool execution
    path in ``agents/agent.py`` bypasses the shared retry-aware middleware chain.

    Args:
        max_retries: Total number of attempts including the first call.  Default 3.
        initial_interval: Seconds to wait before the first retry.  Default 1.0.
        backoff_factor: Multiplier applied to the interval on each retry.  Default 2.0.
        max_interval: Maximum wait time in seconds.  Default 30.0.
        jitter: Randomise delay by ×[0.5, 1.0] to avoid thundering-herd.  Default True.
        retry_on: Tuple of exception types that should trigger a retry.  Default
            ``(Exception,)`` — retries on any exception.
        operation_types: Sequence of operation names (or ``OperationType`` values) to
            apply retry to.  Pass ``None`` to cover *all* operations (advanced use —
            see nested-retry warning above).  Default: ``["tool_call", "memory_access"]``.

    Example::

        from agentbyte.middleware import RetryMiddleware

        agent = Agent(
            client=client,
            middleware=[RetryMiddleware(max_retries=3, initial_interval=0.5, jitter=False)],
        )
    """

    def __init__(
        self,
        max_retries: int = 3,
        initial_interval: float = 1.0,
        backoff_factor: float = 2.0,
        max_interval: float = 30.0,
        jitter: bool = True,
        retry_on: Tuple[Type[BaseException], ...] = (Exception,),
        operation_types: Optional[List[Union[str, OperationType]]] = None,
    ) -> None:
        self.max_retries = max_retries
        self.initial_interval = initial_interval
        self.backoff_factor = backoff_factor
        self.max_interval = max_interval
        self.jitter = jitter
        self.retry_on = retry_on
        # Normalize to frozenset of strings.  None means "use default".
        raw: Optional[Union[frozenset, List[Union[str, OperationType]]]] = (
            _DEFAULT_OPERATION_TYPES if operation_types is None else operation_types
        )
        # Use .value for Enum members so OperationType.TOOL_CALL → "tool_call".
        self._operation_types: Optional[frozenset] = (
            frozenset(op.value if isinstance(op, Enum) else op for op in raw)
            if raw is not None
            else None
        )

    def _compute_delay(self, retry_number: int) -> float:
        """Return wait time in seconds for the given retry number (1-indexed)."""
        delay = self.initial_interval * (self.backoff_factor ** (retry_number - 1))
        delay = min(delay, self.max_interval)
        if self.jitter:
            delay *= random.uniform(0.5, 1.0)
        return delay

    async def process_request(self, context: MiddlewareContext) -> MiddlewareContext:
        """Stamp this instance into context metadata when the operation is applicable."""
        if self._operation_types is None or context.operation in self._operation_types:
            context.metadata[_RETRY_MIDDLEWARE_KEY] = self
        return context

    async def process_response(self, context: MiddlewareContext, result: Any) -> Any:
        return result

    async def process_error(self, context: MiddlewareContext, error: Exception) -> None:
        # Retry logic lives in MiddlewareChain, not here.  This hook runs only
        # after all retry attempts are exhausted, like any other middleware.
        return None
