#!/usr/bin/env python3
# NOTE: PYTHONUNBUFFERED is set in the entrypoint for unbuffered output
#
# An operation that mutates a build by replicating an existing
# EC2 AMI to additional regions, extending the meta.json with
# AMI information about each region replicated to.
#
# Intended to be used in conjunction with buildextend-aws
# to allow the initial upload to a single region for testing
# purposes before replication to a list of regions.

import os
import sys
import json
import argparse
import subprocess

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from cosalib.builds import Builds
from cosalib.cmdlib import write_json

# Parse args and dispatch
parser = argparse.ArgumentParser()
parser.add_argument("--build", help="Build ID",
                    required=True)
parser.add_argument("--name-suffix", help="Suffix for name")
parser.add_argument("--regions", help="EC2 regions, default: all regions",
                    default=[], nargs='+')
parser.add_argument("--log-level", help="ore log level")
args = parser.parse_args()

builds = Builds()
builddir = builds.get_build_dir(args.build)
buildmeta_path = os.path.join(builddir, 'meta.json')
with open(buildmeta_path) as f:
    buildmeta = json.load(f)

base_name = buildmeta['name']
if args.name_suffix:
    ami_name_version = f'{base_name}-{args.name_suffix}-{args.build}'
else:
    ami_name_version = f'{base_name}-{args.build}'


def run_ore():
    if len(buildmeta['amis']) < 1:
        raise SystemExit("buildmeta doesn't contain source AMIs. Run buildextend-aws first")
    if not args.regions:
        args.regions = subprocess.check_output(['ore', 'aws', 'list-regions']).decode().strip().split()
    # only replicate to regions that don't already exist
    existing_regions = [item['name'] for item in buildmeta['amis']]
    duplicates = list(set(args.regions).intersection(existing_regions))
    if len(duplicates) > 0:
        print(f"AMIs already exist in {duplicates} region(s), skipping listed region(s)...")
    region_list = list(set(args.regions) - set(duplicates))
    if len(region_list) == 0:
        print("no new regions detected")
        sys.exit(0)

    source_image = buildmeta['amis'][0]['hvm']
    source_region = buildmeta['amis'][0]['name']
    ore_args = ['ore']
    if args.log_level:
        ore_args.extend(['--log-level', args.log_level])
    ore_args.extend(['aws', 'copy-image', '--image', source_image, '--region', source_region])

    upload_failed_in_region = None
    for upload_region in region_list:
        region_ore_args = ore_args.copy() + [upload_region]
        print("+ {}".format(subprocess.list2cmdline(region_ore_args)))
        try:
            ore_data = json.loads(subprocess.check_output(region_ore_args))
        except subprocess.CalledProcessError:
            upload_failed_in_region = upload_region
            break
        # This matches the Container Linux schema:
        # https://stable.release.core-os.net/amd64-usr/current/coreos_production_ami_all.json
        ami_data = [{'name': region,
                     'hvm': vals['ami'],
                     'snapshot': vals['snapshot']}
                    for region, vals in ore_data.items()]
        buildmeta['amis'].extend(ami_data)

    write_json(buildmeta_path, buildmeta)
    print(f"Updated: {buildmeta_path}")

    if upload_failed_in_region is not None:
        raise Exception(f"Upload failed in {upload_failed_in_region} region")


# Do it!
run_ore()
