from __future__ import annotations

import datetime
import json
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

import boto3

if TYPE_CHECKING:
    from mypy_boto3_sts import STSClient as Boto3STSClient


@dataclass(kw_only=True)
class AssumedRoleCredentials:
    access_key_id: str
    secret_access_key: str
    session_token: str
    expiration: datetime.datetime


class STSClient:
    min_duration = 900

    def __init__(
        self, session: boto3.Session | None = None, **client_kwargs: Any
    ) -> None:
        self._session = session
        self._client_kwargs = client_kwargs
        self._client: Boto3STSClient | None = None

    @property
    def session(self) -> boto3.Session:
        if self._session is None:
            self._session = boto3.Session()
        return self._session

    @property
    def client(self) -> Boto3STSClient:
        if self._client is None:
            self._client = self.session.client("sts", **self._client_kwargs)
        return self._client

    def assume_role(
        self,
        role_arn: str,
        role_session_name: str,
        policy: dict[str, Any],
        duration: int = 900,
    ) -> AssumedRoleCredentials:
        sts_response = self.client.assume_role(
            RoleArn=role_arn,
            RoleSessionName=role_session_name,
            DurationSeconds=max(self.min_duration, duration),
            Policy=json.dumps(policy),
        )

        credentials = sts_response["Credentials"]

        return AssumedRoleCredentials(
            access_key_id=credentials["AccessKeyId"],
            secret_access_key=credentials["SecretAccessKey"],
            session_token=credentials["SessionToken"],
            expiration=credentials["Expiration"],
        )