"""Round-robin and failover API-key pool for provider credential rotation."""

import asyncio
from collections.abc import Sequence

from free_claude_code.providers.failure_policy import (
    is_transient_overload_error,
    retryable_transient_status,
)

_TRANSIENT_ROTATION_STATUSES = frozenset({429, 500, 502, 503, 504})


class ApiKeyPool:
    """Coordinate rotation across one provider's configured credential keys.

    ``acquire`` returns the current key index, rotating only after
    ``per_key_quota`` consecutive acquires on the same key so a balanced burst
    of requests fans out evenly across the pool. ``mark_failure`` jumps
    immediately to the next key when an upstream error is a transient HTTP
    429/5xx or a reported overload, bypassing the per-key budget for that one
    credential. ``advance`` moves the cursor unconditionally and is preserved
    for callers that want strict one-rotate-per-call semantics.
    """

    __slots__ = ("_calls", "_index", "_keys", "_length", "_lock", "_quota")

    def __init__(self, keys: Sequence[str], *, per_key_quota: int = 3) -> None:
        if per_key_quota < 1:
            raise ValueError("per_key_quota must be >= 1")
        self._keys = tuple(keys)
        if not self._keys:
            raise ValueError("ApiKeyPool requires at least one key")
        self._length = len(self._keys)
        self._quota = per_key_quota
        self._index = 0
        self._calls = 0
        self._lock = asyncio.Lock()

    @property
    def keys(self) -> tuple[str, ...]:
        return self._keys

    @property
    def size(self) -> int:
        return self._length

    @property
    def per_key_quota(self) -> int:
        return self._quota

    async def current(self) -> int:
        """Return the current key index without advancing."""
        async with self._lock:
            return self._index

    async def acquire(self) -> int:
        """Return the current key, rotating after ``per_key_quota`` acquires.

        The first ``per_key_quota - 1`` acquires on a key return its index
        unchanged; the ``per_key_quota``-th advance rotates to the next key
        and resets the count for the new key. Concurrent acquires are
        serialized through the event-loop lock so distinct callers never
        observe a stale index or skip a rotation.
        """
        async with self._lock:
            self._calls += 1
            if self._calls >= self._quota:
                self._index = (self._index + 1) % self._length
                self._calls = 0
            return self._index

    async def advance(self) -> int:
        """Advance to the next key (round-robin) and return its new index.

        Legacy unconditional-rotate contract. Resets the per-key acquire
        counter so a subsequent ``acquire`` chain starts fresh on the new key.
        """
        async with self._lock:
            self._index = (self._index + 1) % self._length
            self._calls = 0
            return self._index

    async def mark_failure(self, error: BaseException) -> tuple[int, int] | None:
        """Rotate to the next key when ``error`` warrants credential failover.

        Returns ``(rotated_from_index, rotated_to_index)`` when the error is a
        transient HTTP 429/5xx or a reported upstream overload; otherwise
        returns ``None`` and leaves the pool unchanged. The rotation resets
        the per-key acquire counter so the new key begins a fresh quota
        window.
        """
        if not self._should_rotate(error):
            return None
        async with self._lock:
            previous = self._index
            self._index = (self._index + 1) % self._length
            self._calls = 0
            return previous, self._index

    @staticmethod
    def _should_rotate(error: BaseException) -> bool:
        status = retryable_transient_status(error)
        if status is not None and status in _TRANSIENT_ROTATION_STATUSES:
            return True
        return is_transient_overload_error(error)