#!/bin/bash
#
# Create an AWS AMI from a local disk image (OVA/VMDK/VHD/VHDX/RAW).
#
# Uploads the disk image to an S3 staging bucket, imports it as an EBS
# snapshot, and registers a UEFI-boot AMI.
#
# Prerequisites:
#   - AWS CLI v2 configured with appropriate credentials
#   - jq
#   - Sufficient IAM permissions for S3, EC2 (import-snapshot, register-image)
#   - The vmimport service role configured in the target account (required
#     by ec2:ImportSnapshot; see AWS VM Import/Export prerequisites)
#
# Usage:
#   ./create-ami.sh --disk-image /path/to/image.ova --bucket-name my-staging-bucket
#   ./create-ami.sh --disk-image /path/to/image.vmdk --bucket-name my-bucket --ami-name custom
#   ./create-ami.sh --s3-key image.ova --bucket-name my-bucket --disk-format OVA

set -euo pipefail

# =============================================================================
# Defaults
# =============================================================================
REGION="us-east-1"
DISK_FORMAT=""
AMI_NAME=""
DESCRIPTION=""
KEEP_S3=false
DISK_IMAGE=""
S3_KEY=""
BUCKET_NAME=""
DISK_TYPE="gp3"
DISK_IOPS=3000
DISK_THROUGHPUT=125
POLL_INTERVAL=10

# =============================================================================
# Functions
# =============================================================================

usage() {
    cat <<EOF
Usage: $(basename "$0") [OPTIONS]

Create an AWS AMI (UEFI boot mode) from a local disk image.

Required (one of):
  --disk-image PATH       Path to local disk image file (OVA, VMDK, VHD, VHDX, RAW)
  --s3-key KEY            S3 object key of a disk image already in the bucket (skip upload)

Required:
  --bucket-name NAME      S3 bucket for staging the disk image during import

Optional:
  --region REGION         AWS region (default: us-east-1)
  --ami-name NAME         AMI name (default: derived from filename)
  --description TEXT      AMI description (default: same as AMI name)
  --disk-format FORMAT    Disk format: VMDK, OVA, VHD, VHDX, RAW (default: auto-detect)
  --disk-type TYPE        EBS volume type: gp2, gp3, io1, io2 (default: gp3)
  --disk-iops IOPS        IOPS for gp3/io1/io2 volumes (default: 3000 for gp3)
  --disk-throughput MBps  Throughput in MB/s for gp3 volumes (default: 125)
  --keep-s3               Do not delete the disk image from S3 after AMI creation
  -h, --help              Show this help message

Examples:
  $(basename "$0") --disk-image f5insight.ova --bucket-name my-bucket
  $(basename "$0") --s3-key f5insight.ova --bucket-name my-bucket --disk-format OVA
  $(basename "$0") --disk-image f5insight.ova --bucket-name my-bucket --ami-name custom --disk-type gp3 --disk-iops 4000
EOF
    exit 0
}

die() {
    echo "ERROR: $*" >&2
    exit 1
}

info() {
    echo "==> $*"
}

# Auto-detect disk format from file extension
detect_disk_format() {
    local file="$1"
    local lower
    lower=$(echo "$file" | tr '[:upper:]' '[:lower:]')
    case "$lower" in
        *.vmdk) echo "VMDK" ;;
        *.ova)  echo "OVA" ;;
        *.vhd)  echo "VHD" ;;
        *.vhdx) echo "VHDX" ;;
        *.raw)  echo "RAW" ;;
        *.img)  echo "RAW" ;;
        *)      echo "" ;;
    esac
}

# Derive AMI name from filename
derive_ami_name() {
    local file="$1"
    local basename
    basename=$(basename "$file")
    local lower
    lower=$(echo "$basename" | tr '[:upper:]' '[:lower:]')

    # Strip known extensions
    for ext in .vmdk .ova .vhd .vhdx .raw .img; do
        if [[ "$lower" == *"$ext" ]]; then
            local ext_len=${#ext}
            echo "${basename:0:${#basename}-${ext_len}}"
            return
        fi
    done
    # Fallback: strip last extension
    echo "${basename%.*}"
}

# Replace + with _ for AMI name compatibility
sanitize_ami_name() {
    echo "${1//+/_}"
}

# =============================================================================
# Parse arguments
# =============================================================================
while [[ $# -gt 0 ]]; do
    case "$1" in
        --disk-image)   DISK_IMAGE="$2"; shift 2 ;;
        --s3-key)       S3_KEY="$2"; shift 2 ;;
        --bucket-name)  BUCKET_NAME="$2"; shift 2 ;;
        --region)       REGION="$2"; shift 2 ;;
        --ami-name)     AMI_NAME="$2"; shift 2 ;;
        --description)  DESCRIPTION="$2"; shift 2 ;;
        --disk-format)  DISK_FORMAT="$2"; shift 2 ;;
        --disk-type)    DISK_TYPE="$2"; shift 2 ;;
        --disk-iops)    DISK_IOPS="$2"; shift 2 ;;
        --disk-throughput) DISK_THROUGHPUT="$2"; shift 2 ;;
        --keep-s3)      KEEP_S3=true; shift ;;
        -h|--help)      usage ;;
        *)              die "Unknown option: $1" ;;
    esac
done

# =============================================================================
# Validate inputs
# =============================================================================
[[ -z "$DISK_IMAGE" ]] && [[ -z "$S3_KEY" ]] && die "Must specify either --disk-image (local file) or --s3-key (existing S3 object)"
[[ -z "$BUCKET_NAME" ]] && die "Missing required option: --bucket-name"

# When using a local file, validate it exists
if [[ -n "$DISK_IMAGE" ]] && [[ -z "$S3_KEY" ]]; then
    [[ -f "$DISK_IMAGE" ]] || die "Disk image not found: $DISK_IMAGE"
fi

command -v aws >/dev/null 2>&1 || die "AWS CLI not found. Install it: https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html"
command -v jq >/dev/null 2>&1  || die "jq not found. Install it: https://jqlang.github.io/jq/download/"

# Resolve disk format (use S3_KEY or DISK_IMAGE for detection)
SOURCE_NAME="${S3_KEY:-$DISK_IMAGE}"
if [[ -z "$DISK_FORMAT" ]]; then
    DISK_FORMAT=$(detect_disk_format "$SOURCE_NAME")
    [[ -z "$DISK_FORMAT" ]] && die "Cannot detect disk format from extension. Use --disk-format."
fi
info "Disk format: $DISK_FORMAT"

# Resolve AMI name
if [[ -z "$AMI_NAME" ]]; then
    AMI_NAME=$(derive_ami_name "$SOURCE_NAME")
fi
SANITIZED_NAME=$(sanitize_ami_name "$AMI_NAME")
[[ -z "$DESCRIPTION" ]] && DESCRIPTION="$AMI_NAME"
info "AMI name: $SANITIZED_NAME"

# Determine the S3 key to use for import
if [[ -n "$S3_KEY" ]]; then
    DISK_BASENAME="$S3_KEY"
else
    DISK_BASENAME=$(basename "$DISK_IMAGE")
fi

# =============================================================================
# Upload to S3 (skip when using --s3-key)
# =============================================================================
WE_UPLOADED=false

if [[ -n "$S3_KEY" ]]; then
    info "Using existing S3 object: s3://$BUCKET_NAME/$S3_KEY"
else
    EXISTING_KEY=$(aws s3api head-object \
        --region "$REGION" \
        --bucket "$BUCKET_NAME" \
        --key "$DISK_BASENAME" \
        --query 'ContentLength' \
        --output text 2>/dev/null || echo "")

    if [[ -z "$EXISTING_KEY" ]]; then
        FILE_SIZE=$(stat --printf="%s" "$DISK_IMAGE" 2>/dev/null || stat -f "%z" "$DISK_IMAGE" 2>/dev/null || echo "unknown")
        info "Uploading $DISK_BASENAME ($FILE_SIZE bytes) to s3://$BUCKET_NAME/$DISK_BASENAME"
        aws s3 cp "$DISK_IMAGE" "s3://$BUCKET_NAME/$DISK_BASENAME" --region "$REGION"
        WE_UPLOADED=true
        info "Upload complete"
    else
        info "Disk image '$DISK_BASENAME' already exists in s3://$BUCKET_NAME, skipping upload"
    fi
fi

# =============================================================================
# Import snapshot
# =============================================================================
info "Importing snapshot from s3://$BUCKET_NAME/$DISK_BASENAME (format: $DISK_FORMAT)"

CONTAINER_JSON=$(cat <<EOJSON
{
    "Description": "$DISK_BASENAME",
    "Format": "$DISK_FORMAT",
    "UserBucket": {
        "S3Bucket": "$BUCKET_NAME",
        "S3Key": "$DISK_BASENAME"
    }
}
EOJSON
)

IMPORT_RESPONSE=$(aws ec2 import-snapshot \
    --region "$REGION" \
    --description "$DESCRIPTION" \
    --disk-container "$CONTAINER_JSON" \
    --tag-specifications "ResourceType=import-snapshot-task,Tags=[{Key=Name,Value=${AMI_NAME}}]" \
    --output json)

IMPORT_TASK_ID=$(echo "$IMPORT_RESPONSE" | jq -r '.ImportTaskId')
info "Import task: $IMPORT_TASK_ID"

# Poll until complete
ELAPSED=0
while true; do
    TASK_RESPONSE=$(aws ec2 describe-import-snapshot-tasks \
        --region "$REGION" \
        --import-task-ids "$IMPORT_TASK_ID" \
        --output json)

    STATUS=$(echo "$TASK_RESPONSE" | jq -r '.ImportSnapshotTasks[0].SnapshotTaskDetail.Status')
    PROGRESS=$(echo "$TASK_RESPONSE" | jq -r '.ImportSnapshotTasks[0].SnapshotTaskDetail.Progress // "?"')

    if [[ "$STATUS" == "completed" ]]; then
        SNAPSHOT_ID=$(echo "$TASK_RESPONSE" | jq -r '.ImportSnapshotTasks[0].SnapshotTaskDetail.SnapshotId')
        echo ""
        info "Snapshot import complete: $SNAPSHOT_ID"
        break
    fi

    if [[ "$STATUS" == "deleted" ]]; then
        STATUS_MSG=$(echo "$TASK_RESPONSE" | jq -r '.ImportSnapshotTasks[0].SnapshotTaskDetail.StatusMessage // "unknown error"')
        echo ""
        die "Snapshot import failed: $STATUS_MSG"
    fi

    printf "\rImporting snapshot... %s%% (%ds elapsed)" "$PROGRESS" "$ELAPSED" >&2
    sleep "$POLL_INTERVAL"
    ELAPSED=$((ELAPSED + POLL_INTERVAL))
done

# =============================================================================
# Register AMI (UEFI boot mode)
# =============================================================================
info "Registering AMI: $SANITIZED_NAME"
info "Disk configuration: type=$DISK_TYPE, iops=$DISK_IOPS, throughput=${DISK_THROUGHPUT}MB/s"

# Build EBS configuration based on disk type
EBS_CONFIG="{\"DeleteOnTermination\":true,\"SnapshotId\":\"$SNAPSHOT_ID\",\"VolumeType\":\"$DISK_TYPE\""

if [[ "$DISK_TYPE" == "gp3" ]]; then
    EBS_CONFIG="$EBS_CONFIG,\"Iops\":$DISK_IOPS,\"Throughput\":$DISK_THROUGHPUT"
elif [[ "$DISK_TYPE" =~ ^(io1|io2)$ ]]; then
    EBS_CONFIG="$EBS_CONFIG,\"Iops\":$DISK_IOPS"
fi

EBS_CONFIG="$EBS_CONFIG}"
BLOCK_DEVICE_JSON='[{"DeviceName":"/dev/sda1","Ebs":'"$EBS_CONFIG"'}]'

REGISTER_RESPONSE=$(aws ec2 register-image \
    --region "$REGION" \
    --block-device-mappings "$BLOCK_DEVICE_JSON" \
    --description "$DESCRIPTION" \
    --ena-support \
    --name "$SANITIZED_NAME" \
    --root-device-name "/dev/sda1" \
    --sriov-net-support "simple" \
    --virtualization-type "hvm" \
    --imds-support "v2.0" \
    --boot-mode "uefi" \
    --output json)

AMI_ID=$(echo "$REGISTER_RESPONSE" | jq -r '.ImageId')
info "AMI created: $AMI_ID"

# =============================================================================
# Cleanup S3
# =============================================================================
if [[ "$WE_UPLOADED" == true ]] && [[ "$KEEP_S3" == false ]]; then
    info "Cleaning up s3://$BUCKET_NAME/$DISK_BASENAME"
    aws s3 rm "s3://$BUCKET_NAME/$DISK_BASENAME" --region "$REGION"
fi

echo ""
info "Done. AMI ID: $AMI_ID"