"""S3 utility functions."""
import boto3
import botocore
from lambdacommon.common_config import logger

import config


class S3Util:
    """S3 utility functions."""

    def __init__(self):
        """Initialize S3 client."""
        sts_client = boto3.client('sts')
        assumed_role = sts_client.assume_role(
            RoleArn=config.SHARED_AWS_ROLE_ARN,
            RoleSessionName=config.ROLE_SESSION_NAME
        )
        ACCESS_KEY = assumed_role['Credentials']['AccessKeyId']
        SECRET_KEY = assumed_role['Credentials']['SecretAccessKey']
        SESSION_TOKEN = assumed_role['Credentials']['SessionToken']
        self.s3_client = boto3.resource(
            's3',
            aws_access_key_id=ACCESS_KEY,
            aws_secret_access_key=SECRET_KEY,
            aws_session_token=SESSION_TOKEN,
        )

    def download_audit_log(self, s3_object_path):
        """
        Download an audit log S3 object and save to local temp storage.

        Args:
            s3_object_path (str): path to S3 object
        Returns:
            str: file path on local temp storage where object was saved
        """
        bucket = s3_object_path.split('//')[1].split('/')[0]
        key = '/'.join(s3_object_path.split('//')[1].split('/')[1:])
        file_name = '/tmp/{}'.format(s3_object_path.split('//')[1].split('/')[-1])  # noqa: E501
        try:
            self.s3_client.meta.client.download_file(
                bucket, key, file_name)
            logger.info(f'Downloaded {s3_object_path} to {file_name}')
        except botocore.exceptions.ClientError as error:
            logger.exception('Error downloading S3 object: {}'.format(
                error.response['Error']['Message']))
            raise error

        return file_name