"""Shared middleware building blocks.

This module contains:
- Marker base classes (`Instrumentation`, `Policy`) used to classify middleware in
  `runtime/pipeline.py` (instrumentation vs behavioral policy).
- Helper functions that make middleware implementations work for both sync and async
  executors without duplicating code.
"""

from __future__ import annotations

import inspect
import random
import time
from typing import Any, Awaitable, Callable, Generic, TypeVar, cast

from ..execution import BaseExecRequest, Executor

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


class Instrumentation(Generic[TReq, TOut]):
    """Marker base class for observational middleware.

    Instrumentation middleware should not change semantics. It typically logs, records metrics,
    or emits traces around execution.
    """

    def wrap(self, nxt: Executor[TReq, TOut]) -> Executor[TReq, TOut]:  # pragma: no cover
        raise NotImplementedError


class Policy(Generic[TReq, TOut]):
    """Marker base class for behavioral middleware.

    Policy middleware may change semantics (e.g. retry/backoff, circuit breaking, timeouts).
    """

    def wrap(self, nxt: Executor[TReq, TOut]) -> Executor[TReq, TOut]:  # pragma: no cover
        raise NotImplementedError


def _is_awaitable(x: object) -> bool:
    return inspect.isawaitable(x)


TVal = TypeVar("TVal")


def _ensure_coroutine(x: TVal | Awaitable[TVal]) -> Awaitable[TVal]:
    """Return an awaitable for either sync or async values."""
    if _is_awaitable(x):
        return cast(Awaitable[TVal], x)

    val = cast(TVal, x)

    async def _wrap() -> TVal:
        return val

    return _wrap()


def _call_timed_maybe_async(
    call: Callable[[], TOut],
    *,
    on_success: Callable[[float], None],
    on_error: Callable[[float, Exception], None],
) -> TOut:
    """Call `call()` and run hooks, supporting sync or async return values.

    The `call()` may return either a concrete value (sync executor) or an awaitable
    (async executor). Hooks are triggered after completion with the measured latency.
    """
    start = time.perf_counter()
    try:
        out = call()
    except Exception as e:  # noqa: BLE001
        ms = (time.perf_counter() - start) * 1000
        on_error(ms, e)
        raise

    if _is_awaitable(out):

        async def _run() -> Any:
            try:
                res = await _ensure_coroutine(out)
            except Exception as e:  # noqa: BLE001
                ms = (time.perf_counter() - start) * 1000
                on_error(ms, e)
                raise
            ms = (time.perf_counter() - start) * 1000
            on_success(ms)
            return res

        return cast(TOut, _run())

    ms = (time.perf_counter() - start) * 1000
    on_success(ms)
    return out


def _call_maybe_async(call: Callable[[], TOut], *, on_error: Callable[[Exception], Any]) -> TOut:
    """Call `call()` and error-map, supporting sync or async return values."""
    try:
        out = call()
    except Exception as e:  # noqa: BLE001
        return cast(TOut, on_error(e))

    if _is_awaitable(out):

        async def _run() -> Any:
            try:
                return await _ensure_coroutine(out)
            except Exception as e:  # noqa: BLE001
                return on_error(e)

        return cast(TOut, _run())

    return out


def _retry_delay(policy: Any, attempt: int) -> float:
    """Exponential backoff with optional jitter."""
    delay = min(policy.max_delay_s, policy.base_delay_s * (2 ** (attempt - 1)))
    if getattr(policy, "jitter", None):
        band = max(0.0, float(policy.jitter))
        delay = delay * (1.0 + (random.random() * 2.0 - 1.0) * band)
    return float(delay)

