import re
import os
import logging

import boto3

def list_secrets():
    """
    List Secrets Manager secrets with specific pattern.

    Args:
    prefix (str): Optional. Prefix filter for the secret name.

    Returns:
    list: list of Secrets Manager secret names
    """
    client = boto3.client('secretsmanager')

    env = os.environ.get('ENVIRONMENT', 'dev')
    prefix = f'{env}/direct_delivery/connection_info'

    filters = []
    if prefix:
        filters.append({'Key': 'name', 'Values': [prefix]})

    paginator = client.get_paginator('list_secrets')
    response_iterator = paginator.paginate(Filters=filters, IncludePlannedDeletion=False)

    secret_list = []
    for response in response_iterator:
        if 'SecretList' in response:
            for secret in response['SecretList']:
                secret_list.append(secret['Name'])

    key = rf'{env}/direct_delivery/connection_info.*/release'
    delivery_worker_keys = [secret for secret in secret_list if re.match(key, secret)]

    return delivery_worker_keys

def rollback_secret_to_previous_version(secret_name):
    """
    Rollback an AWS Secrets Manager secret to its previous version.
    
    Parameters:
    - secret_name (str): The name or ARN of the secret.

    Note: This function assumes that the secret versions are labeled with staging labels in an incremental manner.
    """
    client = boto3.client('secretsmanager')

    # List all versions of the secret
    response = client.list_secret_version_ids(SecretId=secret_name)

    versions = response.get('Versions', [])
    if not versions:
        logging.info(f'No versions found for secret: {secret_name}')
        return

    # Sort versions by CreatedDate, most recent first
    sorted_versions = sorted(versions, key=lambda x: x['CreatedDate'], reverse=True)

    # Identify the previous version (the second one in the sorted list)
    if len(sorted_versions) < 2:
        logging.info(f'Only one version found for secret: {secret_name}. Cannot rollback.')
        return

    previous_version = sorted_versions[1]['VersionId']

    # Rollback to the previous version by setting it as the current version
    client.update_secret_version_stage(
        SecretId=secret_name,
        VersionStage='AWSCURRENT', 
        RemoveFromVersionId=sorted_versions[0]['VersionId'], 
        MoveToVersionId=previous_version
    )

    logging.info(f'Rolled back secret {secret_name} to version {previous_version}')

def main():
    """Rollback an AWS Secrets Manager secret to its previous version."""
    secrets = list_secrets()
    for secret in secrets:
        rollback_secret_to_previous_version(secret)


if __name__ == '__main__':
    main()