#!/bin/bash

REGION=$(aws configure get region)
WORKSPACE=/tmp
TEMP_BUCKET_CREATED=false

# Check if user is logged in to AWS
if ! aws sts get-caller-identity >/dev/null 2>&1; then
    echo "Not authenticated with AWS. Please run 'aws configure' or set up your AWS credentials."
    exit 1
fi

# Cleanup function that will always run
on_exit () {
    # Remove temporary VMDK file if it was created
    if [ -n "$vmdk_image_file" ] && [ -e "$WORKSPACE/$vmdk_image_file" ]; then
        rm -f "$WORKSPACE/$vmdk_image_file"
    fi

    # Remove S3 object if it was uploaded
    if [ -n "$S3_BUCKET" ] && [ -n "$vmdk_image_s3_key" ]; then
        aws s3 rm "s3://$S3_BUCKET/$vmdk_image_s3_key" --region "$REGION" >/dev/null 2>&1
    fi

    # Clean up temporary S3 bucket if it was created by this script
    if [ "$TEMP_BUCKET_CREATED" = "true" ] && [ -n "$S3_BUCKET" ]; then
        echo "Cleaning up temporary S3 bucket: $S3_BUCKET"
        # Remove all objects in the bucket first
        aws s3 rm "s3://$S3_BUCKET" --recursive --region "$REGION" >/dev/null 2>&1
        # Delete the bucket
        aws s3api delete-bucket --bucket "$S3_BUCKET" --region "$REGION" >/dev/null 2>&1
    fi
}

get_field () {
	fieldname=$1
	shift
	VAL=$($@|grep $fieldname|sed s/$fieldname//g|sed s/[^a-zA-Z0-9\-]//g)
	echo $VAL
}

function is_tarball()
{
	file $1 | grep -q "tar archive"
}

function is_gzip()
{
	file $1 | grep -q gzip
}

function is_image()
{
	# Check if it's a partition/disk image by file content and has .lh_hdd extension
	file $1 | grep -q partition && [[ $1 == *.lh_hdd ]]
}

function create_temp_s3_bucket()
{
	local timestamp=$(date +%s)
	local bucket_name="lighthouse-bootstrap-temp-${timestamp}"

	# Create bucket with appropriate region constraint
	if [ "$REGION" = "us-east-1" ]; then
		aws s3api create-bucket --bucket "$bucket_name" --region "$REGION" >/dev/null
	else
		aws s3api create-bucket --bucket "$bucket_name" --region "$REGION" --create-bucket-configuration LocationConstraint="$REGION" >/dev/null
	fi

	if [ $? -eq 0 ]; then
		echo "$bucket_name"
		return 0
	fi
}

usage () {
	echo "Usage: $(basename $0) [OPTION]..."
	echo "Creates an AMI from a Lighthouse image file."
	echo
	echo "  -f FILENAME 	Use the specified local file to create the image"
	echo "  -r URI      	Download the image file from the specified URI"
	echo "  -n NAME     	The name to use for generated image. The name cannot contain spaces. (default: Lighthouse)"
	echo "  -s S3_PATH  	Name of S3 bucket to use for temporary storage (optional - will create temp bucket if not provided)"
	echo "  -h          	Display this message"
	echo
}

while getopts "r:f:n:s:" o; do
	case "${o}" in
		f)
			FILENAME=$OPTARG
			;;
		r)
			# If both -r and -f are given, -r will take priority
			DOWNLOAD_LINK=$OPTARG
			;;
		s)
			S3_PATH=$OPTARG
			;;
		n)
			NAME=$OPTARG
			;;
		h)
			usage
			exit 0
			;;
		*)
			usage
			exit 1
			;;
	esac
done

if [ -n "$DOWNLOAD_LINK" ]; then
	echo "Downloading image..."
	wget -q -O /tmp/lh.lh_hdd $DOWNLOAD_LINK
	if [ $? -ne 0 ]; then
		echo "Download failed"
		exit 1
	fi
	FILENAME=/tmp/lh.lh_hdd
fi

if [ -z "$FILENAME" ]; then
	echo "No image path was specified. Either -f or -r options must be provided."
	echo
	usage
	exit 1
fi

if [ ! -f "$FILENAME" ]; then
	echo "Could not locate an image file"
	echo
	usage
	exit 1
fi


if [ -z "$NAME" ]; then
	echo "No name specified. Using the name 'Lighthouse' for the AMI."
	NAME=Lighthouse
else
	# Check if the given name contains spaces
	if [[ "$NAME" =~ [[:space:]] ]]; then
		echo "The AMI name cannot contain spaces. Please provide a valid name."
		echo
		usage
		exit 1
	fi
fi

# Check if S3 bucket exists if provided
if [ -n "$S3_PATH" ]; then
	# Extract bucket name from S3_PATH (everything before the first /)
	S3_BUCKET_TO_CHECK=$(echo "$S3_PATH" | cut -d'/' -f1)
	
	# Check if bucket exists and is accessible
	aws s3api head-bucket --bucket "$S3_BUCKET_TO_CHECK" --region "$REGION" 2>/dev/null
	if [ $? -ne 0 ]; then
		echo "S3 bucket '$S3_BUCKET_TO_CHECK' does not exist or is not accessible"
		exit 1
	fi
fi

# Check image doesn't already exist
IID=$(get_field ImageId aws ec2 describe-images --filter Name=name,Values=${NAME} --region ${REGION})
if [ ! -z "$IID" ]; then
	echo "An image with that name already exists"
	exit 1
fi

##########################
# Convert lh_hdd to vmdk #
##########################

echo "Converting lighthouse image disk (lh_hdd) to VMWare image disk (vmdk)..."

# Extract the tar file to get the image
echo "Extracting image from tar file..."
if is_tarball $FILENAME; then
	tar -xf $FILENAME -C /tmp/
elif is_gzip $FILENAME; then
	tar -zxf $FILENAME -C /tmp/
elif is_image $FILENAME; then
	cp $FILENAME /tmp/
else
	echo "Unsupported file format. Only tar and gzipped tar files are supported."
	exit 1
fi

image_file=$(find /tmp -name "*.lh_hdd" 2>/dev/null | head -1)
if [ -z "$image_file" ]; then
	echo "No .lh_hdd file found in the extracted tar."
	exit 1
fi
vmdk_image_file="${NAME}.vmdk"

# Convert the image to a compressed qcow2 format
image_basename=$(basename "$image_file")
compressed_image_file="${image_basename%.lh_hdd}.qcow2"
if ! qemu-img convert -O qcow2 -c "$WORKSPACE/$image_basename" "$WORKSPACE/$compressed_image_file"; then
	echo "Failed to convert the image to a compressed qcow2 format."
	rm "$image_file"
	exit 1
fi
rm "$WORKSPACE/$image_basename"

# Convert the compressed image to a sparse VMDK format
if ! qemu-img convert -O vmdk -o subformat=streamOptimized "$WORKSPACE/$compressed_image_file" "$WORKSPACE/$vmdk_image_file"; then
	echo "Failed to convert the compressed image to a sparse VMDK format."
	rm "$WORKSPACE/$compressed_image_file"
	exit 1
fi
rm "$WORKSPACE/$compressed_image_file"

#############################
# Create snapshot from vmdk #
#############################

# Set trap to run cleanup on script exit (normal or error)
trap on_exit EXIT

if [ -z "$S3_PATH" ]; then
	echo "No S3 path specified. Creating temporary S3 bucket..."
	S3_PATH=$(create_temp_s3_bucket)
	if [ $? -ne 0 ]; then
		echo "Failed to create temporary S3 bucket"
		exit 1
	fi
	echo "Created temporary S3 bucket: $S3_PATH"
	TEMP_BUCKET_CREATED=true
fi

# make sure our path ends in /lighthouse-aws
S3_PATH="${S3_PATH%/}/lighthouse-aws"
# Extract bucket name and key prefix from S3_PATH
S3_BUCKET=$(echo $S3_PATH | cut -d'/' -f1)
S3_PREFIX=$(echo $S3_PATH | cut -d'/' -f2-)

echo "Uploading VMDK image to S3..."
vmdk_image_s3_key="${S3_PREFIX}/${vmdk_image_file}"
# Upload the VMDK image to S3
if ! aws s3 cp "$WORKSPACE/$vmdk_image_file" "s3://$S3_BUCKET/$vmdk_image_s3_key" --region "$REGION" >/dev/null; then
	echo "Failed to upload the VMDK image to S3."
	exit 1
fi

echo "Importing VMDK image to EC2 as a snapshot..."
# Import the image to EC2 as a snapshot
if ! IMPORT_TASK_ID=$(aws ec2 import-snapshot \
	--disk-container "Format=VMDK,UserBucket={S3Bucket=$S3_BUCKET,S3Key=$vmdk_image_s3_key}" \
	--query 'ImportTaskId' \
	--output text); then
	echo "Failed to import snapshot."
	exit 1
fi

# Monitor the import snapshot task until it is completed
echo "Creating AMI from snapshot: $IMPORT_TASK_ID"
while true; do
	if ! STATUS=$(aws ec2 describe-import-snapshot-tasks --import-task-ids $IMPORT_TASK_ID \
		--query 'ImportSnapshotTasks[0].SnapshotTaskDetail.Status' \
		--output text); then
		echo "Failed to describe import snapshot tasks. $STATUS"
		exit 1
	fi

	if [ "$STATUS" == "completed" ]; then
		break
	elif [ "$STATUS" == "deleted" ] || [ "$STATUS" == "canceled" ]; then
		echo "AMI import task failed: $STATUS"
		exit 1
	fi
	sleep 20
done

if ! SNAPSHOT_ID=$(aws ec2 describe-import-snapshot-tasks --import-task-ids $IMPORT_TASK_ID \
--query 'ImportSnapshotTasks[0].SnapshotTaskDetail.SnapshotId' \
--output text); then
	echo "Failed to describe import snapshot tasks to get SnapshotId. $SNAPSHOT_ID"
	exit 1
fi

############################
# Register the image (AMI) #
############################

echo "Registering AMI..."
if ! AMI_ID=$(aws ec2 register-image --name "$NAME" \
--architecture "x86_64" \
--root-device-name "/dev/sda1" \
--virtualization-type "hvm" \
--ena-support \
--block-device-mappings "DeviceName=/dev/sda1,Ebs={SnapshotId=$SNAPSHOT_ID,VolumeSize=53,VolumeType=gp3}" \
--query 'ImageId' \
--output text); then
	echo "Failed to register the image. $AMI_ID"
	exit 1
fi

echo "AMI created with ID: $AMI_ID"
echo "Done!"