"""Shared retry policy and mixin for LLM clients."""

from __future__ import annotations

import asyncio
import logging
from collections.abc import AsyncGenerator
from typing import Any, Callable, List, Optional

from pydantic import BaseModel, field_validator

logger = logging.getLogger(__name__)


class RetryPolicy(BaseModel):
    """Retry configuration for LLM clients."""

    max_retries: int = 3
    initial_retry_delay: float = 1.0
    max_retry_delay: float = 60.0

    @field_validator("max_retries")
    @classmethod
    def _non_negative_retries(cls, v: int) -> int:
        if v < 0:
            raise ValueError("max_retries must be >= 0")
        return v

    @field_validator("initial_retry_delay", "max_retry_delay")
    @classmethod
    def _no_negative_delays(cls, v: float) -> float:
        if v < 0:
            raise ValueError("delay must be non-negative")
        return v

    def model_post_init(self, __context: Any) -> None:
        if self.max_retry_delay < self.initial_retry_delay:
            raise ValueError(
                "max_retry_delay must be >= initial_retry_delay"
            )


class RetryMixin:
    """Mixin providing unified retry logic for all LLM clients.

    Clients must set ``self.retry_policy`` (a :class:`RetryPolicy` instance)
    before calling these methods.
    """

    retry_policy: RetryPolicy

    async def _retry_with_backoff(
        self,
        func: Callable[..., Any],
        *args: Any,
        _retry_history: Optional[List] = None,
        **kwargs: Any,
    ) -> Any:
        """Execute *func* with exponential back-off on transient errors."""
        from .base import AuthenticationError, InvalidRequestError, RateLimitError
        from .types import ModelClientError
        from ._retry_observability import build_retry_record

        last_exception: Optional[Exception] = None
        current_delay = self.retry_policy.initial_retry_delay

        for attempt in range(self.retry_policy.max_retries + 1):
            try:
                return await func(*args, **kwargs)
            except (RateLimitError, ModelClientError) as e:
                last_exception = e
                if attempt >= self.retry_policy.max_retries:
                    logger.warning(
                        f"Max retries ({self.retry_policy.max_retries}) reached for {func.__name__}"
                    )
                    raise
                if _retry_history is not None:
                    _retry_history.append(
                        build_retry_record(attempt + 1, current_delay, type(e), "retrying")
                    )
                logger.debug(
                    f"Transient error in {func.__name__} "
                    f"(attempt {attempt + 1}/{self.retry_policy.max_retries + 1}): {e}. "
                    f"Retrying in {current_delay:.1f}s"
                )
                await asyncio.sleep(current_delay)
                current_delay = min(current_delay * 2, self.retry_policy.max_retry_delay)
            except (AuthenticationError, InvalidRequestError) as e:
                logger.error(f"Non-transient error in {func.__name__}: {e}")
                raise
            except Exception as e:
                logger.error(f"Unexpected error in {func.__name__}: {e}")
                raise

        raise last_exception if last_exception is not None else ModelClientError("Retry logic error")

    async def _retry_stream_with_backoff(
        self,
        func: Callable[..., AsyncGenerator[Any, None]],
        *args: Any,
        **kwargs: Any,
    ) -> AsyncGenerator[Any, None]:
        """Execute a streaming *func* with exponential back-off on transient errors."""
        from .base import AuthenticationError, InvalidRequestError, RateLimitError
        from .types import ModelClientError

        current_delay = self.retry_policy.initial_retry_delay

        for attempt in range(self.retry_policy.max_retries + 1):
            saw_chunk = False
            try:
                async for chunk in func(*args, **kwargs):
                    saw_chunk = True
                    yield chunk
                return
            except (RateLimitError, ModelClientError) as e:
                if saw_chunk or attempt >= self.retry_policy.max_retries:
                    raise
                logger.debug(
                    f"Transient stream error in {func.__name__} "
                    f"(attempt {attempt + 1}/{self.retry_policy.max_retries + 1}): {e}. "
                    f"Retrying in {current_delay:.1f}s"
                )
                await asyncio.sleep(current_delay)
                current_delay = min(current_delay * 2, self.retry_policy.max_retry_delay)
            except (AuthenticationError, InvalidRequestError):
                raise


__all__ = ["RetryPolicy", "RetryMixin"]
