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

import boto3
from audience_common.aws.utils import get_default_session

from campaigns.utils.asyncio import run_async


@dataclass
class AssumedRoleCredentials:
    access_key_id: str
    secret_access_key: str
    session_token: str
    expiration: datetime


class STSClient:
    def __init__(
        self, region_name: str | None = None, session: boto3.Session | None = None
    ) -> None:
        session = session or get_default_session()
        self.client = session.client("sts", region_name=region_name)

    async def assume_role(
        self,
        role_arn: str,
        role_session_name: str,
        policy: dict[str, Any],
        duration: int = 1800,  # 30 minutes
    ) -> AssumedRoleCredentials:
        sts_response = await run_async(
            self.client.assume_role,
            RoleArn=role_arn,
            RoleSessionName=role_session_name,
            DurationSeconds=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"],
        )