"""STS client wrapper with automatic retry logic."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from alice_cli.aws.session import get_session
from alice_cli.retry import with_aws_retry

if TYPE_CHECKING:
    import boto3


class STSClientWrapper:
    """Wrapper for the boto3 sts client.

    Adds automatic retry logic for transient STS errors, which are rare but
    can occur under high load or during service disruptions.

    Implements the ``STSClient`` protocol.

    Example::

        session = boto3.Session(profile_name="drcc-ai", region_name="us-east-1")
        sts = STSClientWrapper.from_session(session)
        identity = sts.get_caller_identity()
    """

    def __init__(self, client: Any) -> None:
        self._client = client

    @classmethod
    def from_session(cls, session: boto3.Session) -> STSClientWrapper:
        """Create an STSClientWrapper from an existing boto3 Session."""
        return cls(session.client("sts"))

    @classmethod
    def for_profile(cls, profile: str, region: str) -> STSClientWrapper:
        """Create an STSClientWrapper using a cached session for the given profile."""
        session = get_session(profile=profile, region=region)
        return cls.from_session(session)

    @with_aws_retry(max_attempts=3, min_wait=1, max_wait=5)
    def get_caller_identity(self, **kwargs: Any) -> dict[str, Any]:
        """Get the current caller identity with automatic retry.

        Returns:
            Response dict with ``Arn``, ``UserId``, and ``Account`` keys.

        Raises:
            ClientError: If the call fails after retries (e.g. expired credentials).
        """
        return self._client.get_caller_identity(**kwargs)  # type: ignore[no-any-return]
