import boto3
import json
import logging
import os
import config

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def get_ssm_parameter(parameter_name, region):
    """Get the value of an SSM parameter."""
    ssm_client = boto3.client("ssm", region_name=region)
    response = ssm_client.get_parameter(Name=parameter_name, WithDecryption=False)
    return response["Parameter"]["Value"]


def parse_account_id_and_role(parameter_value):
    """Parse the account ID and role from an SSM parameter value."""
    data = json.loads(parameter_value)
    account_id = data.get("account_id")
    role = data.get("role")
    return account_id, role


def assume_role(account_id, role_name, session_name):
    """Assume an IAM role in another account."""
    role_arn = f"arn:aws:iam::{account_id}:role/{role_name}"
    sts_client = boto3.client("sts")
    assume_role_response = sts_client.assume_role(
        RoleArn=role_arn, RoleSessionName=session_name
    )
    return assume_role_response["Credentials"]


def list_s3_objects(bucket_name, prefix, region, credentials):
    """List all objects in an S3 bucket with the given prefix."""
    s3_client = boto3.client(
        "s3",
        region_name=region,
        aws_access_key_id=credentials["AccessKeyId"],
        aws_secret_access_key=credentials["SecretAccessKey"],
        aws_session_token=credentials["SessionToken"],
    )
    response = s3_client.list_objects_v2(Bucket=bucket_name, Prefix=prefix)
    return [obj["Key"] for obj in response.get("Contents", [])]


def store_secrets_to_file(bucket_name, s3_keys, file_path, region, credentials):
    """Store secrets from S3 to a file and list them.
    The secrets are stored in a dictionary with the secret name as the key
    Account ID is stripped to match target account secret name space."""
    s3_client = boto3.client(
        "s3",
        region_name=region,
        aws_access_key_id=credentials["AccessKeyId"],
        aws_secret_access_key=credentials["SecretAccessKey"],
        aws_session_token=credentials["SessionToken"],
    )
    secrets = {}
    for s3_key in s3_keys:
        secret_name = s3_key.replace(config.BUCKET_SECRETS_ACCOUNT + "/", "").rstrip(":")
        s3_client.download_file(bucket_name, s3_key, file_path)
        with open(file_path, "r") as file:
            content = file.read()
            try:
                secret_value = json.loads(content)
                secrets[secret_name] = secret_value
            except json.JSONDecodeError:
                secrets[secret_name] = content
    with open(file_path, "w") as file:
        json.dump(secrets, file)

    # Log the secrets
    logging.info("Secrets stored in the file:")
    for secret_name, secret_value in secrets.items():
        logging.info(f"{secret_name}: {secret_value}")


def update_secret(secret_name, secret_value, region, credentials):
    """Update a secret in AWS Secrets Manager."""
    secrets_client = boto3.client(
        "secretsmanager",
        region_name=region,
        aws_access_key_id=credentials["AccessKeyId"],
        aws_secret_access_key=credentials["SecretAccessKey"],
        aws_session_token=credentials["SessionToken"],
    )
    try:
        # Update the secret
        response = secrets_client.update_secret(
            SecretId=secret_name, SecretString=secret_value
        )
        logger.info(f"Successfully updated secret: {secret_name}")
    except Exception as e:
        logger.error(f"Error updating secret {secret_name}: {e}")


def main():
    bucket_name = config.BACKUP_SECRETS_BUCKET
    prefix = config.BUCKET_SECRETS_PREFIX
    file_path = config.SECRET_FILE
    region = config.AWS_REGION

    try:
        # Get source account ID and role
        source_parameter_name = config.AWS_ACCOUNT_IDS_SSM + config.SOURCE_AWS_ACCOUNT
        source_parameter_value = get_ssm_parameter(source_parameter_name, region)
        source_account_id, source_role_name = parse_account_id_and_role(
            source_parameter_value
        )

        # Assume the source role
        current_user = boto3.client("iam").get_user()["User"]["UserName"]
        source_credentials = assume_role(
            source_account_id, source_role_name, current_user
        )

        # List and store secrets from S3
        s3_keys = list_s3_objects(bucket_name, prefix, region, source_credentials)
        store_secrets_to_file(
            bucket_name, s3_keys, file_path, region, source_credentials
        )

        # Verify if the file was created successfully
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"The file {file_path} was not created.")

        # Get destination account ID and role
        destination_parameter_name = (
            config.AWS_ACCOUNT_IDS_SSM + config.DESTINATION_AWS_ACCOUNT
        )
        destination_parameter_value = get_ssm_parameter(
            destination_parameter_name, region
        )
        destination_account_id, destination_role_name = parse_account_id_and_role(
            destination_parameter_value
        )

        # Assume the destination role
        destination_credentials = assume_role(
            destination_account_id, destination_role_name, current_user
        )

        # Update each secret in AWS Secrets Manager
        with open(file_path, "r") as file:
            secrets = json.load(file)
        for secret_name, secret_value in secrets.items():
            update_secret(secret_name, secret_value, region, destination_credentials)

    except Exception as e:
        logging.error(f"An error occurred: {e}")


if __name__ == "__main__":
    main()