#!/usr/bin/env bash
# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

set -o errexit
set -o nounset
set -o pipefail


EKSD_RELEASE_BRANCH="${1?Specify first argument - EKS-D release branch}"
AMI_MANIFEST_OUTPUT="${2?Specify second argument - AMI manifest output file from packer}"
IMAGE_FORMAT="${3?Specify third argument - Format for exported image \(vmdk\|raw\|vhd\)}"
S3_DST_EXPORT_PATH="${4?Specify fourth argument - Destination S3 path}"
REPLICAS="${5?Specify fourth argument - Comma separated list of s3 destinations for exported image copies}"

SCRIPT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd -P)"
source "${SCRIPT_ROOT}/build/lib/common.sh"

ARTIFACT_ID=$(cat $AMI_MANIFEST_OUTPUT | jq -r '.builds[0].artifact_id')
ARTIFACT_ID_SPLIT=(${ARTIFACT_ID//:/ })
AMI_ID=${ARTIFACT_ID_SPLIT[1]}

REPLICAS_SPLIT=(${REPLICAS//,/ })

DST_BUCKET_PATH=${S3_DST_EXPORT_PATH#"s3://"}
DST_BUCKET_NAME=${DST_BUCKET_PATH%%/*}
DST_PATH=${DST_BUCKET_PATH#"$DST_BUCKET_NAME/"}

EKSD_RELEASE=$(build::eksd_releases::get_eksd_release_name $EKSD_RELEASE_BRANCH)
EXPORTED_IMAGE_PREFIX="$DST_PATH/$EKSD_RELEASE-"

EXPORT_RESPONSE=$(aws ec2 export-image --disk-image-format $IMAGE_FORMAT --s3-export-location S3Bucket=$DST_BUCKET_NAME,S3Prefix=$EXPORTED_IMAGE_PREFIX --image-id $AMI_ID)
echo $EXPORT_RESPONSE

EXPORT_TASK_ID=$(echo $EXPORT_RESPONSE | jq -r '.ExportImageTaskId')
EXPORTED_IMAGE_LOCATION="s3://${DST_BUCKET_NAME}/${EXPORTED_IMAGE_PREFIX}${EXPORT_TASK_ID}.${IMAGE_FORMAT}"

FINAL_STATUSES=(completed deleted)
STATUS=$(echo $EXPORT_RESPONSE | jq -r '.Status')
STATUS_MESSAGE=$(echo $EXPORT_RESPONSE | jq -r '.StatusMessage')
PROGRESS=$(echo $EXPORT_RESPONSE | jq -r '.Progress')

until [[ "${FINAL_STATUSES[*]}" =~ "${STATUS}" ]]; do
  echo "Image import is $STATUS: $STATUS_MESSAGE $PROGRESS%"
  sleep 30

  DESCRIBE_RESPONSE=$(aws ec2 describe-export-image-tasks --export-image-task-ids $EXPORT_TASK_ID)
  STATUS=$(echo $DESCRIBE_RESPONSE | jq -r '.ExportImageTasks[0].Status')
  STATUS_MESSAGE=$(echo $DESCRIBE_RESPONSE | jq -r '.ExportImageTasks[0].StatusMessage')
  PROGRESS=$(echo $DESCRIBE_RESPONSE | jq -r '.ExportImageTasks[0].Progress')
done

if [[ "$STATUS" != "completed" ]]; then
    echo "Image import failed: $STATUS - $STATUS_MESSAGE"
    exit 1
fi

echo "Image import for ami $AMI_ID succeeded"

for dst in "${REPLICAS_SPLIT[@]}"
do
  echo "Copying exported image to $dst"
  aws s3 cp --no-progress --acl public-read $EXPORTED_IMAGE_LOCATION $dst 
done