# -*- coding: utf-8 -*-
"""Retry wrapper for ChatModelBase instances.

Transparently retries LLM API calls on transient errors (rate-limit,
timeout, connection) with configurable exponential back-off.

Concurrency and rate-limit control (LLMRateLimiter):
- A global semaphore caps the number of concurrent in-flight LLM calls,
  preventing a burst of requests from hammering the upstream API.
- When a 429 is received every concurrent caller is paused for the same
  duration (plus per-caller jitter) before re-trying, eliminating the
  thundering-herd problem where multiple callers retry at the same instant.

Semaphore ownership rules:
- Non-streaming: __call__'s finally block always releases the slot
  (owns_semaphore stays True throughout).
- Streaming: ownership transfers to _consume_stream_with_slot the moment
  __call__ returns the generator.  owns_semaphore is set to False before
  the return so __call__'s finally skips the release.
  _consume_stream_with_slot releases after the first chunk arrives.
- Cancellation safety: the boolean flag `acquired` tracks whether the
  semaphore slot has actually been taken; the final block only releases
  when acquired is True, preventing a spurious release on CancelledError.
"""

from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass
from typing import Any, AsyncGenerator

from agentscope.model import ChatModelBase
from agentscope.model._model_response import ChatResponse
from qwenpaw.exceptions import (
    RateLimitExceededException,
)

from ..constant import (
    LLM_ACQUIRE_TIMEOUT,
    LLM_BACKOFF_BASE,
    LLM_BACKOFF_CAP,
    LLM_MAX_CONCURRENT,
    LLM_MAX_RETRIES,
    LLM_MAX_QPM,
    LLM_RATE_LIMIT_JITTER,
    LLM_RATE_LIMIT_PAUSE,
)
from .model_capability_cache import get_capability_cache
from .rate_limiter import LLMRateLimiter, get_rate_limiter

logger = logging.getLogger(__name__)

RETRYABLE_STATUS_CODES = {429, 500, 502, 503, 504, 529}

_openai_retryable: tuple[type[Exception], ...] | None = None
_anthropic_retryable: tuple[type[Exception], ...] | None = None
_httpx_retryable: tuple[type[Exception], ...] | None = None


class _AcquireTimeoutError(RateLimitExceededException):
    """Raised when ``limiter.acquire()`` times out internally.

    Distinct from a real API 429 so the retry loop can identify it via
    ``isinstance`` and raise immediately without calling
    ``report_rate_limit()`` or attempting another retry.
    """


@dataclass(frozen=True, slots=True)
class RetryConfig:
    """Retry policy for transient LLM API failures."""

    enabled: bool = LLM_MAX_RETRIES > 0
    max_retries: int = max(LLM_MAX_RETRIES, 1)
    backoff_base: float = LLM_BACKOFF_BASE
    backoff_cap: float = LLM_BACKOFF_CAP


@dataclass(frozen=True, slots=True)
class RateLimitConfig:
    """Rate-limiting policy for LLM calls.

    Controls the global LLMRateLimiter singleton that caps concurrency and
    coordinates pauses when a 429 is received.  The singleton is initialised
    on the *first* call; subsequent callers share the same instance.

    Attributes:
        max_concurrent: Maximum concurrent in-flight LLM calls.
        max_qpm: Maximum queries per minute (sliding window). 0 = disabled.
        pause_seconds: Global pause duration (s) on a 429 response.
        jitter_range: Random jitter (s) added on top of the pause.
        acquire_timeout: Max seconds to wait for a slot before raising.
    """

    max_concurrent: int = LLM_MAX_CONCURRENT
    max_qpm: int = LLM_MAX_QPM
    pause_seconds: float = LLM_RATE_LIMIT_PAUSE
    jitter_range: float = LLM_RATE_LIMIT_JITTER
    acquire_timeout: float = LLM_ACQUIRE_TIMEOUT


def _get_openai_retryable() -> tuple[type[Exception], ...]:
    global _openai_retryable
    if _openai_retryable is None:
        try:
            import openai

            _openai_retryable = tuple(
                cls
                for cls in (
                    openai.RateLimitError,
                    openai.APITimeoutError,
                    openai.APIConnectionError,
                    getattr(openai, "InternalServerError", None),
                )
                if cls is not None
            )
        except ImportError:
            _openai_retryable = ()
    return _openai_retryable


def _get_anthropic_retryable() -> tuple[type[Exception], ...]:
    global _anthropic_retryable
    if _anthropic_retryable is None:
        try:
            import anthropic

            _anthropic_retryable = (
                anthropic.RateLimitError,
                anthropic.APITimeoutError,
                anthropic.APIConnectionError,
            )
        except ImportError:
            _anthropic_retryable = ()
    return _anthropic_retryable


def _get_httpx_retryable() -> tuple[type[Exception], ...]:
    global _httpx_retryable
    if _httpx_retryable is None:
        try:
            import httpx

            _httpx_retryable = (
                httpx.RemoteProtocolError,
                httpx.TimeoutException,
            )
        except ImportError:
            _httpx_retryable = ()
    return _httpx_retryable


def _extract_status_code(exc: Exception) -> int | None:
    """Best-effort HTTP status extraction from SDK exceptions.

    Streaming SSE errors are raised as plain ``openai.APIError`` without a
    ``status_code`` attribute; the gateway status is often only present in
    ``body`` (e.g. ``{"status_code": 502, "error": {...}}``).
    """
    status = getattr(exc, "status_code", None)
    if status is not None:
        try:
            return int(status)
        except (TypeError, ValueError):
            pass

    body = getattr(exc, "body", None)
    if not isinstance(body, dict):
        return None

    for container in (body, body.get("error")):
        if not isinstance(container, dict):
            continue
        raw = container.get("status_code")
        if raw is None:
            raw = container.get("code")
        if raw is None:
            continue
        try:
            return int(raw)
        except (TypeError, ValueError):
            continue

    return None


def _is_retryable(exc: Exception) -> bool:
    """Return *True* if *exc* should trigger a retry."""
    retryable = (
        _get_openai_retryable()
        + _get_anthropic_retryable()
        + _get_httpx_retryable()
    )
    if retryable and isinstance(exc, retryable):
        return True

    status = _extract_status_code(exc)
    if status is not None and status in RETRYABLE_STATUS_CODES:
        return True

    return False


def _is_rate_limit(exc: Exception) -> bool:
    """Return *True* if *exc* is specifically a 429 rate-limit error."""
    return _extract_status_code(exc) == 429


def _is_missing_reasoning_content_error(exc: Exception) -> bool:
    """Return *True* if *exc* is a 400 about missing ``reasoning_content``.

    DeepSeek (and compatible providers) require every assistant message to
    carry ``reasoning_content`` when thinking mode is active.  When the
    conversation history was produced by a non-reasoning model, these
    fields are absent and the API rejects the request with a 400.
    """
    if _extract_status_code(exc) != 400:
        return False
    return "reasoning_content" in str(exc)


def _inject_reasoning_content(
    args: tuple,
    kwargs: dict[str, Any],
) -> bool:
    """Add ``reasoning_content = " "`` to assistant messages that lack it.

    Modifies the formatted message dicts **in-place** so the subsequent
    retry sees the updated values.  Returns *True* when at least one
    message was patched.
    """
    messages: list[dict] | None = kwargs.get("messages")
    if messages is None and args:
        candidate = args[0]
        if isinstance(candidate, list):
            messages = candidate

    if not messages:
        return False

    modified = False
    for msg in messages:
        if (
            isinstance(msg, dict)
            and msg.get("role") == "assistant"
            and "reasoning_content" not in msg
        ):
            msg["reasoning_content"] = " "
            modified = True

    return modified


def _extract_retry_after(exc: Exception) -> float | None:
    """Parse the Retry-After header value (in seconds) from an exception.

    Handles both OpenAI and Anthropic SDK exception shapes, which expose
    headers either directly on the exception or on an attached response object.
    """
    headers = getattr(exc, "headers", None) or getattr(
        getattr(exc, "response", None),
        "headers",
        None,
    )
    if headers:
        raw = headers.get("Retry-After") or headers.get("retry-after")
        if raw:
            try:
                return float(raw)
            except (TypeError, ValueError):
                pass
    return None


def _normalize_retry_config(retry_config: RetryConfig | None) -> RetryConfig:
    """Normalize externally supplied retry config into safe bounds."""
    if retry_config is None:
        return RetryConfig()
    normalized_backoff_base = max(0.1, retry_config.backoff_base)
    normalized_backoff_cap = max(
        0.5,
        retry_config.backoff_cap,
        normalized_backoff_base,
    )
    return RetryConfig(
        enabled=retry_config.enabled,
        max_retries=max(1, retry_config.max_retries),
        backoff_base=normalized_backoff_base,
        backoff_cap=normalized_backoff_cap,
    )


def _normalize_rate_limit_config(
    cfg: RateLimitConfig | None,
) -> RateLimitConfig:
    """Normalize externally supplied rate-limit config into safe bounds."""
    if cfg is None:
        return RateLimitConfig()
    return RateLimitConfig(
        max_concurrent=max(1, cfg.max_concurrent),
        max_qpm=max(0, cfg.max_qpm),
        pause_seconds=max(1.0, cfg.pause_seconds),
        jitter_range=max(0.0, cfg.jitter_range),
        acquire_timeout=max(10.0, cfg.acquire_timeout),
    )


def _compute_backoff(attempt: int, retry_config: RetryConfig) -> float:
    """Exponential back-off: base * 2^(attempt-1), capped."""
    return min(
        retry_config.backoff_cap,
        retry_config.backoff_base * (2 ** max(0, attempt - 1)),
    )


class RetryChatModel(ChatModelBase):
    """Transparent retry wrapper around any :class:`ChatModelBase`.

    The wrapper delegates every call to the underlying *inner* model and
    retries on transient errors with exponential back-off.  Streaming
    responses are also covered: if the stream fails mid-consumption the
    entire request is retried from scratch.

    A global LLMRateLimiter is consulted on every call to cap concurrency and
    to coordinate a shared pause across all callers when a 429 is received.
    """

    def __init__(
        self,
        inner: ChatModelBase,
        retry_config: RetryConfig | None = None,
        rate_limit_config: RateLimitConfig | None = None,
    ) -> None:
        # agentscope 2.0 ChatModelBase requires credential/model/parameters;
        # forward the inner wrapper's own values so attribute access stays
        # transparent.
        super().__init__(
            credential=getattr(inner, "credential", None),
            model=getattr(inner, "model", "unknown"),
            parameters=getattr(inner, "parameters", None)
            or ChatModelBase.Parameters(),
            stream=getattr(inner, "stream", True),
            context_size=getattr(inner, "context_size", 32768),
        )
        self._inner = inner
        self._retry_config = _normalize_retry_config(retry_config)
        self._rate_limit_config = _normalize_rate_limit_config(
            rate_limit_config,
        )

    # Expose the real model's class so that formatter mapping keeps working
    # when code inspects ``model.__class__`` after wrapping.
    @property
    def inner_class(self) -> type:
        return self._inner.__class__

    @property
    def model_key(self) -> str:
        """Stable key for the underlying model: ``provider_id:model_name``."""
        provider_id = getattr(self._inner, "_provider_id", None)
        name = self._inner.model
        return f"{provider_id}:{name}" if provider_id else name

    @staticmethod
    async def _handle_rate_limit_exc(
        exc: Exception,
        limiter: LLMRateLimiter,
    ) -> None:
        """Inspect *exc* and update the rate limiter accordingly.

        - Internal acquire timeout (``_AcquireTimeoutError``): re-raise as-is;
          no report, no retry.
        - Retryable API 429 with Retry-After > ``MAX_PAUSE_SECONDS``: re-raise
          immediately — retrying after the capped pause would just get another
          429 (e.g. Anthropic FreeUsageLimitError with Retry-After: 51496 s).
        - Normal 429: call ``report_rate_limit()`` to set the per-model pause.
        """
        if isinstance(exc, _AcquireTimeoutError):
            raise exc
        if _is_retryable(exc) and _is_rate_limit(exc):
            retry_after = _extract_retry_after(exc)
            if (
                retry_after is not None
                and retry_after > LLMRateLimiter.MAX_PAUSE_SECONDS
            ):
                raise exc
            await limiter.report_rate_limit(retry_after)

    async def _consume_stream_with_slot(
        self,
        stream: AsyncGenerator[ChatResponse, None],
        limiter: LLMRateLimiter,
        acquired_at: float,
    ) -> AsyncGenerator[ChatResponse, None]:
        """Yield all chunks from *stream*, managing the semaphore slot
        lifecycle.

        Releases the semaphore slot after the first chunk arrives — once the
        API starts streaming the request has been accepted and will not be
        rate-limited mid-flight, so holding the slot for the full streaming
        duration would unnecessarily starve other callers.

        Always closes *stream* on completion or error.  Any exception raised
        during iteration propagates to the caller's ``async for`` loop
        (i.e. _wrap_stream), which handles retry decisions.  The exception
        does not propagate to the final consumer unless all retries are
        exhausted.

        Args:
            acquired_at: Timestamp from ``limiter.acquire()``, forwarded to
                ``on_success()`` so only stale pauses are cleared.
        """
        first_chunk = True
        try:
            async for chunk in stream:
                if first_chunk:
                    first_chunk = False
                    # return the slot once the API starts delivering
                    limiter.release()
                    # streaming success: clear any stale 429 pause so
                    # subsequent callers (including user chats) are not
                    # held back by a pause set by a background task.
                    await limiter.on_success(acquired_at)
                yield chunk
        finally:
            await stream.aclose()
            if first_chunk:
                # Stream failed before producing any chunk;
                # slot not yet released.
                limiter.release()

    async def generate_structured_output(
        self,
        *args: Any,
        **kwargs: Any,
    ) -> Any:
        return await self._inner.generate_structured_output(*args, **kwargs)

    async def __call__(
        self,
        *args: Any,
        **kwargs: Any,
    ) -> ChatResponse | AsyncGenerator[ChatResponse, None]:
        cache = get_capability_cache()
        key = self.model_key

        if cache.get(key, "needs_reasoning_content", False):
            _inject_reasoning_content(args, kwargs)

        # Each model gets its own rate limiter keyed by
        # "provider_id:model_name" so that a 429 on one model (e.g. from a
        # dream/cron task) cannot stall user chats on a different provider.
        limiter = await get_rate_limiter(
            limiter_key=self.model_key,
            max_concurrent=self._rate_limit_config.max_concurrent,
            max_qpm=self._rate_limit_config.max_qpm,
            default_pause_seconds=self._rate_limit_config.pause_seconds,
            jitter_range=self._rate_limit_config.jitter_range,
        )

        retries = (
            self._retry_config.max_retries if self._retry_config.enabled else 0
        )
        attempts = retries + 1
        last_exc: Exception | None = None

        for attempt in range(1, attempts + 1):
            # Acquire a semaphore slot, with a timeout to prevent
            # indefinite blocking. `acquired` tracks whether the slot was
            # taken so the final block can skip the release on
            # CancelledError (slot was never acquired).
            acquired = False
            owns_semaphore = True
            acquired_at: float = 0.0
            try:
                try:
                    acquired_at = await asyncio.wait_for(
                        limiter.acquire(),
                        timeout=self._rate_limit_config.acquire_timeout,
                    )
                    acquired = True
                except asyncio.TimeoutError as exc:
                    # Internal acquire timeout — NOT an API 429.
                    # _AcquireTimeoutError is a typed subclass so the outer
                    # handler can use isinstance() instead of a sentinel attr.
                    raise _AcquireTimeoutError(
                        operation="LLM execution",
                        retry_after=int(
                            self._rate_limit_config.acquire_timeout,
                        ),
                        details={
                            "reason": "Timed out waiting for execution slot",
                        },
                    ) from exc

                try:
                    result = await self._inner(*args, **kwargs)
                except Exception as inner_exc:
                    if not (
                        _is_missing_reasoning_content_error(inner_exc)
                        and _inject_reasoning_content(args, kwargs)
                    ):
                        raise
                    cache.learn(key, "needs_reasoning_content", True)
                    logger.warning(
                        "Thinking-mode model requires reasoning_content "
                        "on every assistant message. Injecting empty "
                        "values and retrying (learned for future calls).",
                    )
                    result = await self._inner(*args, **kwargs)

                if isinstance(result, AsyncGenerator):
                    # Transfer semaphore ownership to _wrap_stream, which uses
                    # _consume_stream_with_slot internally and handles
                    # retries on stream failure.
                    owns_semaphore = False
                    return self._wrap_stream(
                        result,
                        args,
                        kwargs,
                        attempt,
                        attempts,
                        limiter,
                        acquired_at,
                    )

                # Non-streaming success: clear any stale rate-limit pause so
                # subsequent callers are not held back by a pause set by an
                # unrelated background task (e.g. dream/cron 429).
                await limiter.on_success(acquired_at)
                return result

            except Exception as exc:
                last_exc = exc
                await self._handle_rate_limit_exc(exc, limiter)

                if not _is_retryable(exc) or attempt >= attempts:
                    raise

                delay = _compute_backoff(attempt, self._retry_config)
                logger.warning(
                    "LLM call failed (attempt %d/%d): %s. "
                    "Retrying in %.1fs ...",
                    attempt,
                    attempts,
                    exc,
                    delay,
                )
                await asyncio.sleep(delay)

            finally:
                if owns_semaphore and acquired:
                    limiter.release()

        # Should be unreachable, but satisfies the type-checker.
        raise last_exc  # type: ignore[misc]

    # pylint: disable=too-many-branches
    async def _wrap_stream(
        self,
        stream: AsyncGenerator[ChatResponse, None],
        call_args: tuple,
        call_kwargs: dict,
        current_attempt: int,
        max_attempts: int,
        limiter: LLMRateLimiter,
        acquired_at: float = 0.0,
    ) -> AsyncGenerator[ChatResponse, None]:
        """Yield chunks from *stream*; on transient failure, retry the full
        request and yield from the new stream instead.

        Args:
            acquired_at: Timestamp from ``limiter.acquire()``, forwarded to
                ``on_success()`` so stale pauses are cleared but fresh ones
                (set by a concurrent 429 after this call acquired) are kept.
        """
        attempt = current_attempt
        pending_stream: AsyncGenerator[ChatResponse, None] | None = stream
        pending_acquired_at = acquired_at
        reasoning_injected = False

        while True:
            try:
                if pending_stream is not None:
                    async for chunk in self._consume_stream_with_slot(
                        pending_stream,
                        limiter,
                        pending_acquired_at,
                    ):
                        yield chunk
                    return  # stream completed without error

                acquired = False
                owns_semaphore = True
                retry_acquired_at: float = 0.0
                try:
                    try:
                        retry_acquired_at = await asyncio.wait_for(
                            limiter.acquire(),
                            timeout=self._rate_limit_config.acquire_timeout,
                        )
                        acquired = True
                    except asyncio.TimeoutError as exc:
                        raise _AcquireTimeoutError(
                            operation="LLM execution (stream retry)",
                            retry_after=int(
                                self._rate_limit_config.acquire_timeout,
                            ),
                            details={
                                "reason": (
                                    "Timed out waiting for execution slot"
                                ),
                            },
                        ) from exc

                    result = await self._inner(*call_args, **call_kwargs)

                    if isinstance(result, AsyncGenerator):
                        owns_semaphore = False
                        pending_stream = result
                        pending_acquired_at = retry_acquired_at
                        continue

                    yield result
                    return
                finally:
                    if owns_semaphore and acquired:
                        limiter.release()

            except Exception as retry_exc:
                pending_stream = None
                if (
                    not reasoning_injected
                    and _is_missing_reasoning_content_error(retry_exc)
                    and _inject_reasoning_content(call_args, call_kwargs)
                ):
                    reasoning_injected = True
                    get_capability_cache().learn(
                        self.model_key,
                        "needs_reasoning_content",
                        True,
                    )
                    logger.warning(
                        "Thinking-mode stream requires reasoning_content "
                        "on every assistant message. Injecting empty "
                        "values and retrying (learned for future calls).",
                    )
                    continue

                if _is_retryable(retry_exc) and _is_rate_limit(retry_exc):
                    await limiter.report_rate_limit(
                        _extract_retry_after(retry_exc),
                    )

                if not _is_retryable(retry_exc) or attempt >= max_attempts:
                    raise

                retry_delay = _compute_backoff(attempt, self._retry_config)
                logger.warning(
                    "LLM stream failed (attempt %d/%d): %s. "
                    "Retrying in %.1fs ...",
                    attempt,
                    max_attempts,
                    retry_exc,
                    retry_delay,
                )
                await asyncio.sleep(retry_delay)
                attempt += 1
