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 get_s3_zone_file_content(bucket_name, key, region, credentials):
    """Retrive restore file in an S3 bucket with the given key."""
    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.get_object(Bucket=bucket_name, Key=key)
    content = response['Body'].read().decode('utf-8')
    return content


def find_hosted_zone_id(domain_name, region, credentials):
    """Find the hosted zone ID for a given domain name in Route 53."""
    route53_client = boto3.client(
        "route53",
        region_name=region,
        aws_access_key_id=credentials["AccessKeyId"],
        aws_secret_access_key=credentials["SecretAccessKey"],
        aws_session_token=credentials["SessionToken"],
    )
    hosted_zones = route53_client.list_hosted_zones_by_name(DNSName=domain_name)
    
    for zone in hosted_zones["HostedZones"]:
        if zone["Name"].rstrip('.') == domain_name:
            return zone["Id"].split('/')[-1]  # Extract only the hosted zone ID
    raise Exception(f"Hosted zone for {domain_name} not found.")


def import_zone_file_content(zone_id, zone_file_content, region, credentials):
    """Import zone file content into a Route 53 hosted zone by creating/updating records."""
    route53_client = boto3.client(
        "route53",
        region_name=region,
        aws_access_key_id=credentials["AccessKeyId"],
        aws_secret_access_key=credentials["SecretAccessKey"],
        aws_session_token=credentials["SessionToken"],
    )
    records = parse_zone_file(zone_file_content)

    changes = []
    for record in records:
        changes.append({
            "Action": "UPSERT",
            "ResourceRecordSet": record
        })

    response = route53_client.change_resource_record_sets(
        HostedZoneId=zone_id,
        ChangeBatch={
            "Comment": "Automated zone file import",
            "Changes": changes
        }
    )
    logger.info("Zone file imported successfully.")
    return response

def parse_zone_file(zone_file_content):
    """Parse zone file content into Route 53 compatible format."""
    records = []
    for line in zone_file_content.strip().splitlines():
        parts = line.split()
        
        # Skip lines that don’t have the expected format
        if len(parts) < 5:
            logger.warning(f"Skipping unrecognized line format: {line}")
            continue

        try:
            name = parts[0]
            ttl = int(parts[1])
            record_type = parts[3]
            value = parts[4]

            if record_type == "TXT":
                value = '"' + value.strip('"') + '"'  # Ensure TXT values are quoted

            record = {
                "Name": name,
                "Type": record_type,
                "TTL": ttl,
                "ResourceRecords": [{"Value": value}]
            }
            records.append(record)
        except ValueError as e:
            logger.error(f"Error parsing line: {line} - {e}")
            continue
    return records


def main():
    bucket_name = config.BACKUP_SECRETS_BUCKET
    key = config.RECORDS_FILE_S3_KEY
    region = config.AWS_REGION
    domain_name = config.DOMAIN_NAME

    # 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
    )

    # Retrieve the zone file content from S3
    zone_file_content = get_s3_zone_file_content(bucket_name, key, region, source_credentials)
    
     # 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
        )
    # Find the hosted zone ID for the domain
    zone_id = find_hosted_zone_id(domain_name, region, destination_credentials)

    # Import the zone file content into Route 53
    import_response = import_zone_file_content(zone_id, zone_file_content, region, destination_credentials)
    print("Import completed:", import_response)

if __name__ == "__main__":
    main()