"""Module for restoring RDS snapshots."""
import datetime
import json
import logging
import sys

import boto3
from boto3 import Session
from botocore.config import Config
from botocore.credentials import RefreshableCredentials
from botocore.session import get_session
import sentry_sdk
from sentry_sdk.integrations.aws_lambda import AwsLambdaIntegration

import config
from common.logic import database
from common.logic import sanitiser

if config.SENTRY_DSN:
    sentry_sdk.init(
        dsn=config.SENTRY_DSN,
        environment=config.ENVIRONMENT,
        integrations=[AwsLambdaIntegration(timeout_warning=True)]
    )

logger = logging.getLogger()
logger.setLevel(config.LOGGER_LEVEL)
stream_handler = logging.StreamHandler(sys.stdout)
stream_handler.setLevel(logging.DEBUG)
logger.handlers = [stream_handler]


def assume_target_account_role():
    """
    Assume IAM role in target account.

    Args:
    account_id (str): AWS account ID

    Returns:
    dict: dict of AWS credentials
    """
    client = boto3.client('sts')
    assume_role_response = client.assume_role(
        RoleArn=f'arn:aws:iam::{config.TARGET_ACCOUNT_ID}:role/{config.TARGET_ACCOUNT_ROLE_NAME}', # noqa
        RoleSessionName=config.TARGET_ACCOUNT_ROLE_SESSION_NAME,
        ExternalId=config.EXTERNAL_ID,
        DurationSeconds=config.ROLE_DURATION_SECONDS
    )
    credentials = assume_role_response['Credentials']

    return {
        'access_key': credentials['AccessKeyId'],
        'secret_key': credentials['SecretAccessKey'],
        'token': credentials['SessionToken'],
        'expiry_time': credentials['Expiration'].isoformat()
    }


def get_rds_client():
    """
    Return a new RDS client.

    The client will automatically refresh the credentials when they expire.
    """
    session_credentials = RefreshableCredentials.create_from_metadata(
        metadata=assume_target_account_role(),
        refresh_using=assume_target_account_role,
        method='sts-assume-role',
    )

    session = get_session()
    session._credentials = session_credentials
    refreshable_session = Session(botocore_session=session)
    rds_client = refreshable_session.client(
        'rds', region_name=config.AWS_DEFAULT_REGION)

    return rds_client


def get_lambda_client():
    """Return a new Lambda client."""
    credentials = assume_target_account_role()
    lambda_client = boto3.client(
        'lambda',
        region_name=config.AWS_DEFAULT_REGION,
        aws_access_key_id=credentials['access_key'],
        aws_secret_access_key=credentials['secret_key'],
        aws_session_token=credentials['token'],
        config=Config(
            retries={'max_attempts': 0},
            read_timeout=config.LAMBDA_TIMEOUT,
            connect_timeout=config.LAMBDA_TIMEOUT
        )
    )

    return lambda_client


def copy_snapshot(rds_client):
    """Copy a shared snapshot to a local snapshot."""
    local_snapshot_name = config.SNAPSHOT_NAME.removesuffix('-shared')
    logger.info(
        f'Copying shared snapshot {config.SNAPSHOT_NAME} '
        f'to local snapshot {local_snapshot_name}')

    database.copy_snapshot(
        config.SNAPSHOT_NAME,
        config.SOURCE_ACCOUNT_ID,
        config.DB_TYPE,
        local_snapshot_name,
        'destination',
        config.RESTORE_KMS_KEY_ID,
        config.AWS_DEFAULT_REGION,
        rds_client,
        additional_tags=get_snapshot_tags())

    return local_snapshot_name


def get_snapshot_tags():
    """Return any additional tags to add to snapshots."""
    return [
        {'Key': 'execution_id', 'Value': config.EXECUTION_ID},
    ]


def restore(db_name, source_db_name, db_type, snapshot_name, rds_client):
    """Restore an RDS snapshot to an existing cluster or clone the database."""
    timestamp = datetime.datetime.now().strftime('%Y%m%d%H%M%S')
    intermediate_db_name = f'{db_name}-tmp-{timestamp}'
    old_db_name = f'{db_name}-old'
    existing_db_config = database.get_config(
        db_name, db_type, rds_client)
    logger.info(f'Existing DB configuration: {existing_db_config}')

    refresh_complete = False
    try:
        clone = config.CLONE_RESTORE and db_type.lower() == 'cluster'
        if clone:
            logger.info(f'Cloning database {db_name} to {intermediate_db_name}.')

            source_db_arn = (
                f'arn:aws:rds:{config.AWS_DEFAULT_REGION}:'
                f'{config.SOURCE_ACCOUNT_ID}:'
                f'cluster:{source_db_name}'
            )
            database.clone_database(
                source_db_arn, intermediate_db_name, db_type, existing_db_config,
                rds_client)
        else:
            logger.info(
                f'Restoring snapshot to intermediate {intermediate_db_name}.')

            snapshot_type = 'cluster-snapshot' if db_type.lower() == 'cluster' else 'snapshot' # noqa
            snapshot_arn = (
                f'arn:aws:rds:{config.AWS_DEFAULT_REGION}:'
                f'{config.SOURCE_ACCOUNT_ID}:'
                f'{snapshot_type}:{snapshot_name}'
            )
            database.restore_database(
                snapshot_arn, intermediate_db_name, db_type, existing_db_config,
                rds_client)

        logger.info(f'Sanitising data in {intermediate_db_name}')
        sanitiser.sanitise(
            config.SANITISE_DATA_FUNCTION_NAME,
            db_name,
            intermediate_db_name,
            db_type,
            f'scripts/{db_name}',
            existing_db_config['Engine'],
            get_lambda_client()
        )

        logger.info(f'Renaming existing instance {db_name} to {old_db_name}.')
        database.rename_database(db_name, old_db_name, db_type, rds_client)

        logger.info(
            f'Renaming intermediate instance {intermediate_db_name} to {db_name}.')
        database.rename_database(
            intermediate_db_name, db_name, db_type, rds_client)

        logger.info('Configuring event subscriptions')
        database.configure_event_subscriptions(
            db_name, db_type, existing_db_config, rds_client)

        refresh_complete = True

    finally:
        if refresh_complete:
            logger.info(f'Deleting old instance {old_db_name} if it exists.')
            database.delete_database_if_exists(old_db_name, db_type, rds_client, wait=False)
        else:
            # If the refresh failed after the live database was renamed to
            # -old but before the intermediate took its place, -old holds the
            # only copy of the previous data, so never delete it on failure;
            # leave it for manual recovery.
            logger.warning(
                f'Refresh failed; leaving {old_db_name} in place if it exists.')

        logger.info(f'Deleting intermediate instance {intermediate_db_name} if it exists.')
        database.delete_database_if_exists(intermediate_db_name, db_type, rds_client, wait=False)


def snapshot_copy_required(
        source_account_id, target_account_id, db_type):
    """
    Return true if a snapshot copy is required to perform the restore.

    This is only required when performing a cross-account restore of a
    snapshot, and only for standalone databases.
    """
    is_cross_account = source_account_id and source_account_id != target_account_id # noqa
    return is_cross_account and db_type == 'standalone'


def main():
    """Task entry point."""
    sfn_client = boto3.client(
        'stepfunctions', region_name=config.AWS_DEFAULT_REGION)
    try:
        rds_client = get_rds_client()

        snapshot_name = config.SNAPSHOT_NAME
        if snapshot_copy_required(
                config.SOURCE_ACCOUNT_ID, config.TARGET_ACCOUNT_ID,
                config.DB_TYPE):
            snapshot_name = copy_snapshot(rds_client)

        restore(
            db_name=config.DB_NAME,
            source_db_name=config.SOURCE_DB_NAME,
            db_type=config.DB_TYPE,
            snapshot_name=snapshot_name,
            rds_client=rds_client
        )

        if config.TASK_TOKEN:
            sfn_client.send_task_success(
                taskToken=config.TASK_TOKEN,
                output=json.dumps(config.DB_NAME)
            )
    except Exception as error:
        logger.exception(str(error))
        if config.TASK_TOKEN:
            sfn_client.send_task_failure(
                taskToken=config.TASK_TOKEN,
                error='500',
                cause=repr(error)
            )
        raise error


if __name__ == '__main__':
    main()