"""Route53 zones backup."""
import logging
from datetime import datetime
from datetime import timezone


import boto3
import sentry_sdk


import config
from connectors import datadog

sentry_sdk.init(
    dsn=config.SENTRY_DSN,
    traces_sample_rate=1.0,
)

logger = logging.getLogger()
logger.setLevel(config.LOGGER_LEVEL)


def assume_source_account_role(account_id):
    """
    Assume IAM role in source 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::{account_id}:role/{config.CROSS_ACCOUNT_BACKUP_ROLE_NAME}',
        RoleSessionName=config.SERVICE_NAME,
        DurationSeconds=3600
    )
    credentials = assume_role_response['Credentials']

    return credentials


def fetch_zone_records(source_account_id):
    """
    List Route53 zone info for a given account.

    Args:
    source_account_id (str): AWS account ID

    Returns:
    list: list of Route53 zone info dictionaries
    """
    credentials = assume_source_account_role(source_account_id)
    client = boto3.client(
        'route53',
        region_name=config.AWS_REGION,
        aws_access_key_id=credentials['AccessKeyId'],
        aws_secret_access_key=credentials['SecretAccessKey'],
        aws_session_token=credentials['SessionToken']
    )
    zone_info_list = []
    paginator = client.get_paginator('list_hosted_zones')
    response_iterator = paginator.paginate()
    for response in response_iterator:
        if 'HostedZones' in response:
            for zone in response['HostedZones']:
                zone_id = zone['Id']
                zone_name = zone['Name'].strip('.')

                client.get_hosted_zone(Id=zone['Id'])
                zone_info = {
                    'Name': zone_name,
                    'Config': zone['Config'],
                    'ResourceRecordSets': []
                }

                paginator = client.get_paginator('list_resource_record_sets')

                for record_set in paginator.paginate(HostedZoneId=zone_id):
                    zone_info['ResourceRecordSets'].extend(record_set['ResourceRecordSets'])

                zone_info_list.append(zone_info)

    return zone_info_list


def convert_to_bind_format(zone_info_list):
    """
    Convert Route53 zone info to BIND format, excluding SOA and NS records with the same name as the zone.

    Args:
    zone_info_list (list): A list of Route53 zone info dictionaries

    Returns:
    dict: A dictionary where keys are zone names and values are BIND format data for each zone
    """
    bind_format_data = {}

    for zone_info in zone_info_list:
        bind_format = []
        zone_name = zone_info['Name']

        bind_format.append(f'$ORIGIN {zone_name}.\n')
        bind_format.append(f"$TTL {zone_info['Config'].get('TTL', 3600)}\n")  # Default TTL if not provided

        for record_set in zone_info['ResourceRecordSets']:
            record_name = record_set['Name']
            record_type = record_set['Type']
            ttl = record_set.get('TTL', 300)

            if record_type == 'SOA' or (record_type == 'NS' and record_name == f"{zone_name}."):
                continue

            if 'ResourceRecords' in record_set:
                for record in record_set['ResourceRecords']:
                    bind_format.append(f"{record_name} {ttl} IN {record_type} {record['Value']}\n")

            if 'AliasTarget' in record_set:
                alias_target = record_set['AliasTarget']['DNSName']
                bind_format.append(f"{record_name} {ttl} IN {record_type} {alias_target} ; ALIAS\n")

        bind_format_data[zone_name] = ''.join(bind_format)

    return bind_format_data


def put_hosted_zone_data_to_s3(client, bucket_name, zone_name, bind_data, source_account_id):
    """
    Backup Route53 zones to s3 bucket in BIND format.

    Args:
    client (obj): a boto3 client
    bucket_name (str): Backup bucket name
    zone_name (str): S3 key for the zone location
    bind_data (str): BIND format data
    """
    timestamp = datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ')
    client.put_object(
        Bucket=bucket_name,
        Key=f'{source_account_id}/{zone_name}/backup_{timestamp}.txt',  # Save as .txt in BIND format
        Body=bind_data,
    )


def main():
    """Backup Route53 zones to s3 bucket."""
    s3_client = boto3.client('s3', region_name=config.BACKUP_AWS_REGION)
    bucket_name = config.BUCKET_NAME
    account_id = config.ACCOUNT_ID_TO_BACKUP
    zone_info_list = fetch_zone_records(account_id)

    try:
        bind_format_data = convert_to_bind_format(zone_info_list)

        for zone_name, bind_data in bind_format_data.items():
            put_hosted_zone_data_to_s3(s3_client, bucket_name, zone_name, bind_data, account_id)
            logging.info(f'Backing up {account_id}/{zone_name} Route53 zone to s3 in BIND format')

        logging.info(f'Route53 backup for account {account_id} is completed')
        datadog.publish_metric('route53_backup.success', 1)

    except Exception as error:
        logger.exception(str(error))
        datadog.publish_metric('route53_backup.error', 1)
        raise error


if __name__ == '__main__':
    main()