"""Retry/backoff policy middleware.

This middleware is *behavioral policy*:
- It can re-execute the underlying executor on retryable failures.
- It adds delays between attempts using exponential backoff with optional jitter.

Important separation of concerns:
- This module implements the *retry mechanism*.
- Retryability *classification* (what counts as transient) lives in
  `relationalai.runtime.retry_classification` and is injected via `is_retryable`.
"""

from __future__ import annotations

import asyncio
import logging
import time
from dataclasses import dataclass
from typing import Any, Awaitable, Callable, Generic, TypeVar, cast

from ..execution import BaseExecRequest, Executor
from .base import _ensure_coroutine, _is_awaitable, _retry_delay
from .metrics import MetricsSink, NoOpMetricsSink

TReq = TypeVar("TReq", bound=BaseExecRequest)
TOut = TypeVar("TOut")


@dataclass(frozen=True)
class RetryPolicy:
    """Retry/backoff configuration.

    - `max_attempts`: total attempts including the first attempt.
    - delays are exponential (`base_delay_s * 2^(attempt-1)`) capped at `max_delay_s`.
    - `jitter`: randomizes delay by +/- jitter% to avoid thundering herds.
    """

    max_attempts: int = 1
    base_delay_s: float = 0.1
    max_delay_s: float = 1.0
    jitter: float = 0.25


class RetryMiddleware(Generic[TReq, TOut]):
    """Retry wrapper for sync or async execution.

    This middleware supports both:
    - sync executors (where `nxt.execute(req)` raises immediately)
    - async executors (where `nxt.execute(req)` returns an awaitable)

    Metrics:
    - increments `<kind>.retry` for attempts > 1 with `meta['attempt']=attempt`.
    """

    def __init__(
        self,
        *,
        policy: RetryPolicy,
        is_retryable: Callable[[Exception, TReq], bool],
        logger: logging.Logger | None = None,
        metrics: MetricsSink | None = None,
        enabled: bool = True,
    ) -> None:
        self._policy = policy
        self._is_retryable = is_retryable
        self._log = logger or logging.getLogger(__name__)
        self._metrics = metrics or NoOpMetricsSink()
        self._enabled = enabled

    def wrap(self, nxt: Executor[TReq, TOut]) -> Executor[TReq, TOut]:
        policy = self._policy
        is_retryable = self._is_retryable
        log = self._log
        metrics = self._metrics
        enabled = self._enabled

        class _Wrapped:
            def execute(self, req: TReq) -> TOut:
                if not enabled:
                    return nxt.execute(req)

                attempts = max(1, int(policy.max_attempts))
                try:
                    first = nxt.execute(req)
                except Exception as e:  # noqa: BLE001
                    # Sync failure (first attempt raised) -> sync retry loop, including the first failure.
                    last: Exception | None = e
                    for attempt in range(1, attempts + 1):
                        if attempt > 1:
                            # Keep retry counters separate from attempt/success/error metrics.
                            metrics.inc(
                                operation=req.operation,
                                kind=f"{req.kind}.retry",
                                meta={**req.meta, "attempt": attempt},
                            )
                        try:
                            if attempt == 1:
                                raise e
                            return nxt.execute(req)
                        except Exception as e2:  # noqa: BLE001
                            last = e2
                            if attempt < attempts and is_retryable(e2, req):
                                delay = _retry_delay(policy, attempt)
                                log.warning(
                                    "exec_retry kind=%s op=%s attempt=%s/%s delay=%.2f err=%s",
                                    req.kind,
                                    req.operation,
                                    attempt,
                                    attempts,
                                    delay,
                                    str(e2)[:200],
                                )
                                if delay > 0:
                                    time.sleep(delay)
                                continue
                            raise
                    assert last is not None
                    raise last

                if not _is_awaitable(first):
                    # Sync success on first try.
                    return first

                async def _run_async() -> Any:
                    last: Exception | None = None
                    for attempt in range(1, attempts + 1):
                        if attempt > 1:
                            metrics.inc(
                                operation=req.operation,
                                kind=f"{req.kind}.retry",
                                meta={**req.meta, "attempt": attempt},
                            )
                        try:
                            if attempt == 1:
                                # Reuse the already-created awaitable from the first attempt.
                                return await _ensure_coroutine(cast(Awaitable[Any], first))
                            return await _ensure_coroutine(nxt.execute(req))
                        except Exception as e:  # noqa: BLE001
                            last = e
                            if attempt < attempts and is_retryable(e, req):
                                delay = _retry_delay(policy, attempt)
                                log.warning(
                                    "exec_retry kind=%s op=%s attempt=%s/%s delay=%.2f err=%s",
                                    req.kind,
                                    req.operation,
                                    attempt,
                                    attempts,
                                    delay,
                                    str(e)[:200],
                                )
                                if delay > 0:
                                    await asyncio.sleep(delay)
                                continue
                            raise
                    assert last is not None
                    raise last

                # Async-first attempt: return a coroutine that performs the retry loop.
                return cast(Any, _run_async())

        return _Wrapped()

