@PYTHON@

import re
import yaml
import boto3


# All Alpine AMIs should match this regex if they're valid
AMI_RE = re.compile("^Alpine-(\d+\.\d+)(?:-r(\d+))?-Hardened-EC2")


# Load current AMI version from config
with open("alpine-ami.yaml") as fp:
    ami_cfg = yaml.full_load(fp)["variables"]
    current = (float(ami_cfg["alpine_release"]), int(ami_cfg["ami_release"]))


# Fetch all matching AMIs
amis = []

for region in boto3.session.Session().get_available_regions("ec2"):
    ec2 = boto3.client("ec2", region_name=region)

    for image in ec2.describe_images(Owners=["self"])["Images"]:
        match = AMI_RE.match(image["Name"])
        if not match:
            continue
        
        os_rel, ami_rel = match.groups()
        amis.append((
            region, image["ImageId"],
            image["BlockDeviceMappings"][0]["Ebs"]["SnapshotId"],
            float(os_rel), int(ami_rel) if ami_rel else 0))


# Determine the set to discard based region and version
ok_regions = set()
discards = []

# Cluster candidates by region/version pair, newest in a region first.
# This should result in the first match for a region always being the newest
# AMI for that region and all subsequent matches in the region being old.
# Even so we must keep track of regions with current images on the off-chance
# that a region only has old images. In that case we want to preserve the old
# images till we can publish new ones manually so users can still launch
# Alpine systems without interruption.
candidates = sorted(amis, key=lambda i: (i[0], (i[1], i[3])), reverse=True)

for ami in candidates:
    (region, ami, snapshot), version = ami[:3], ami[3:]

    if version > current:
        print("{} has AMI '{}' newer than current".format(region, ami))
        continue
    elif version == current:
        ok_regions.add(region)
        continue
    elif version < current and region in ok_regions:
        discards.append((region, ami, snapshot))
    else:
        print("Not discarding old image in {}".format(region))
        continue


# Scrub the old ones
for region, image, snapshot in discards:
    print("Removing image '{}', snapshot '{}' in {}".format(
        image, snapshot, region))

    ec2 = boto3.client("ec2", region_name=region)
    ec2.deregister_image(ImageId=image)
    ec2.delete_snapshot(SnapshotId=snapshot)