#!/usr/bin/env bash

VERSION="$1"
EPOCH=${EPOCH:-0}

RLVER=$VERSION

# Source common variables
# shellcheck disable=SC2046,1091,1090
source "$(dirname "${BASH_SOURCE[0]}")/common"

usage() {
	echo "usage: $0 VERSION ($0 8.7)"
}

aws() {
	command aws --region us-east-1 --profile resf-ami --output text $@
}

DATE=${DATE:-$(date +%Y%m%d)}

if [[ -z $VERSION ]]; then
	usage
	exit 1
fi

download() {
	curl -o "$qcow" "https://dl.rockylinux.org/stg/rocky/$VERSION/images/$arch/$qcow"
	return $?
}

exists() {
	aws s3 ls $2/$1 &>/dev/null
	return $?
}

upload() {
	if exists $raw $BUCKET/$DATE; then
		echo "Found existing upload in $BUCKET"
		return 0
	fi
	aws s3 cp $raw $BUCKET/$DATE/
}

convert() {
	qemu-img convert -f qcow2 -O raw "$qcow" $raw
	return $?
}

BUCKET="s3://resf-prod-import-use1"

write-container-file() {
	printf '{"Description": "%s", "Format": "raw", "Url": "%s/%s"}\n' $name "$BUCKET/$DATE" $raw >$json
}

download-convert-upload() {
	if ! [[ -f $qcow ]]; then
		echo "Downloading $qcow from cdn"
		download || exit
	fi

	if ! [[ -f ${raw} ]]; then
		echo "Converting qcow2 $qcow to RAW $raw"
		convert || exit
	fi

	echo "Uploading raw image back to s3"
	if ! upload; then
		echo "Failed to upload"
		exit
	fi
}

begin-job() {
	import_task_id="$(aws ec2 import-snapshot --disk-container "file://$PWD/$json" --query 'ImportTaskId')"
	if [[ -z $import_task_id ]]; then
		echo "Failed to import $json"
		return 1
	fi
	echo $import_task_id
	return 0
}

is-snapshot-imported() {
	snapshot_id=$(aws ec2 describe-import-snapshot-tasks --query 'ImportSnapshotTasks[].SnapshotTaskDetail.SnapshotId[]' --import-task-ids $1)
	if [[ -z "$snapshot_id" ]]; then
		return 1
	fi
	return 0
}

register-image() {
	# Given a snapshot id, register the image
	name=$1
	snapshot_id=$2

	case $(awk -F'.' '{print $NF}' <<<"$name") in
	x86_64)
		arch=x86_64
		;;
	aarch64)
		arch=arm64
		;;
	esac

	ami_id=$(aws --query "ImageId" ec2 register-image --name "$name" --description "$name" --block-device-mappings DeviceName="/dev/sda1",Ebs={SnapshotId="$snapshot_id"} --root-device-name "/dev/sda1" --virtualization-type hvm --architecture $arch --ena-support)

	if [[ -z "$ami_id" ]]; then
		return 1
	fi
	return 0
}

tag-resources() {
	local resources="$1"
	local tags="$2"
	if [[ -z $resources || -z $tags ]]; then
		echo "Need to provide tags and resources to tag"
		return 1
	fi
	aws ec2 create-tags --resources $resources --tags $tags
}

image-exists() {
	local RESF_AMI_ACCOUNT_ID=792107900819
	local query="$(printf 'Images[?Name==`%s`].[ImageId,Name][]' "${1}")"
	mapfile -t res < <(aws ec2 describe-images --owners $RESF_AMI_ACCOUNT_ID --query "${query}" 2>/dev/null)
	res=($res)
	if [[ ${#res[@]} -eq 0 ]]; then
		# Skip empty results
		return 1 #not found
	fi
	id=${res[0]//\"/}
	name=${res[@]/$id/}
	found_image_id=$id
	return 0 # found
}
snapshot-exists() {
	local RESF_AMI_ACCOUNT_ID=792107900819
	local filter="$(printf 'Name=tag:Name,Values=%s' "${1}")"
	local query='Snapshots[].[SnapshotId][]'
	mapfile -t res < <(aws ec2 describe-snapshots --owner-ids $RESF_AMI_ACCOUNT_ID --filter "${filter}" --query "${query}" 2>/dev/null)
	res=($res)
	if [[ ${#res[@]} -eq 0 ]]; then
		# Skip empty results
		return 1 #not found
	fi
	id=${res[0]//\"/}
	found_snapshot_id=$id
	return 0 # found
}

declare -A import_jobs
declare -A snapshot_ids
declare -A ami_ids

TARGET_ARCHES=(x86_64 aarch64)
TARGET_VARIANTS=(Base LVM)

for variant in "${TARGET_VARIANTS[@]}"; do
	for arch in "${TARGET_ARCHES[@]}"; do
		latest=$(printf "Rocky-%s-EC2-%s.latest.%s" "$VERSION" $variant $arch)
		name=$(printf "Rocky-%s-EC2-%s-%s-%s.%s.%s" "$VERSION" $variant $REVISION $DATE $EPOCH $arch)
		qcow=${name}.qcow2
		raw=${name}.raw
		json=${name}.json

		if image-exists $name; then
			echo "Found existing AMI in us-east-1. Skipping. ($found_image_id,$name)"
			continue
		fi

		if snapshot-exists $name; then
			# If the snapshot exists, we can skip the import task and just do the image registration
			echo "Found existing snapshot: ($found_snapshot_id,$name)"
			snapshot_ids[$name]="${found_snapshot_id}"
			continue
		fi

		# Download latest artifacts from CDN, convert from qcow2 to raw, and upload to the proper bucket
		echo "Downloading/converting artifacts for $name"
		download-convert-upload

		echo "Writing disk container json file"
		write-container-file

		jobid=$(begin-job)
		echo "Beginning snapshot import task with id $jobid"
		import_jobs[$name]=$jobid
	done
done

# wait for all import jobs to complete, then tag the resultant images
finished=false
while ! $finished; do
	for name in "${!import_jobs[@]}"; do
		import_task_id="${import_jobs[${name}]}"
		if ! is-snapshot-imported $import_task_id; then
			echo "Snapshot for $import_task_id ($name) is not yet finished"
			continue
		fi

		# await finalization
		sleep 2

		if [[ -z $snapshot_id ]]; then
			echo "Snapshot ID is null.. exiting"
			exit 2
		fi

		echo "Tagging snapshot with name"
		tag-resources $snapshot_id "Key=Name,Value=$name"

		unset import_jobs[${name}]
		snapshot_ids[${name}]=$snapshot_id
	done
	# Check if we're done, if so, great!
	if [[ ${#import_jobs[@]} -gt 0 ]]; then
		echo "Sleeping for 1m"
		sleep 1m
		continue
	else
		finished=true
		break
	fi
done

finished=false
while ! $finished; do
	for name in "${!snapshot_ids[@]}"; do
		# If the snapshot is imported, turn it into an AMI
		snapshot_id="${snapshot_ids[${name}]}"

		echo "Creating AMI from snapshot $snapshot_id ($name)"
		if ! register-image $name $snapshot_id; then
			echo "ERROR: Failed to create image for $name with snapshot id $snapshot_id"
			continue
		fi

		echo "Tagging AMI - Name=$name"
		tag-resources $ami_id "Key=Name,Value=$name"

		if [[ -z $ami_id ]]; then
			echo "AMI ID is null. continuing..."
			continue
		fi

		unset snapshot_ids[${name}]
		ami_ids[$name]=$ami_id
	done
	if [[ ${#snapshot_ids[@]} -gt 0 ]]; then
		sleep 1m
		continue
	else
		finished=true
		break
	fi
done

res=""
for name in "${!ami_ids[@]}"; do
	ami_id="${ami_ids[${name}]}"
	res="${res}\n$(printf '%s\t%s\n' $name $ami_id)"
done

printf "$res\n"