import os
import logging
from typing import Optional, Dict, Any
import boto3
from botocore.exceptions import ClientError
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


class STS:
    """
    STS helper for AWS Security Token Service operations.
    """

    def __init__(self, region_name: Optional[str] = None):
        self.region = region_name or os.getenv("AWS_DEFAULT_REGION", "us-east-1")
        self.client = boto3.client("sts", region_name=self.region)

    def assume_role(self, params: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """
        Assume an IAM role.
        :param params: Dictionary containing AssumeRole parameters
        :return: boto3 response dict or None on error
        """
        try:
            response = self.client.assume_role(**params)
            return response
        except ClientError as e:
            logger.error("STS ClientError in assume_role: %s", e)
            return None
        except Exception as e:
            logger.exception("Unexpected error in assume_role: %s", e)
            return None

    def get_identity(self) -> Optional[Dict[str, Any]]:
        """
        Get the caller identity.
        :return: boto3 response dict or None on error
        """
        try:
            response = self.client.get_caller_identity()
            return response
        except ClientError as e:
            logger.error("STS ClientError in get_identity: %s", e)
            return None
        except Exception as e:
            logger.exception("Unexpected error in get_identity: %s", e)
            return None


