#!/usr/bin/env bash
set -euo pipefail

aws_region="${AWS_REGION:-us-west-2}"
system="aarch64-linux"
builder_instance_type="t4g.large"
name="nixos-zfs-ec2"
keep_builder=false
extra_tofu_args=()

usage() {
  cat >&2 <<'EOF'
Usage: publish-ami.sh [options] [-- <extra tofu apply args>]

Options:
  --system <aarch64-linux|x86_64-linux>
                              Target NixOS system (default: aarch64-linux)
  --aws-region <region>       AWS region (default: AWS_REGION or us-west-2)
  --builder-instance-type <t> Temporary builder EC2 type (default: t4g.large)
  --name <name>               AMI name prefix
  --keep-builder              Keep temporary builder resources after publish
  -h, --help                  Show this help

Any arguments after `--` are passed through to `tofu apply`.
Use that for values like:

  -- -var base_nixos_ami_id=ami-... -var subnet_id=subnet-...
EOF
}

while [ "$#" -gt 0 ]; do
  case "$1" in
    --system)
      system="$2"
      shift 2
      ;;
    --aws-region)
      aws_region="$2"
      shift 2
      ;;
    --builder-instance-type)
      builder_instance_type="$2"
      shift 2
      ;;
    --name)
      name="$2"
      shift 2
      ;;
    --keep-builder)
      keep_builder=true
      shift
      ;;
    --)
      shift
      extra_tofu_args=("$@")
      break
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "Unknown argument: $1" >&2
      usage
      exit 1
      ;;
  esac
done

case "$system" in
  aarch64-linux)
    architecture="arm64"
    ;;
  x86_64-linux)
    architecture="x86_64"
    ;;
  *)
    echo "Unsupported system: $system" >&2
    exit 1
    ;;
esac

require() {
  if ! command -v "$1" >/dev/null 2>&1; then
    echo "Missing required command: $1" >&2
    exit 1
  fi
}

require aws
require git
require jq
require sha256sum
require tofu

repo_root="$(git rev-parse --show-toplevel)"
if ! git diff --quiet || ! git diff --cached --quiet; then
  echo "The AMI builder git tree must be clean before publishing." >&2
  exit 1
fi

builder_git_rev="$(git rev-parse HEAD)"
flake_lock_rev="$(jq -r '.nodes.nixpkgs.locked.rev // .nodes.nixpkgs.locked.narHash // "unknown"' "$repo_root/flake.lock")"
input_hash="$(printf '%s\n%s\n%s\n' "$builder_git_rev" "$flake_lock_rev" "$system" | sha256sum | awk '{ print $1 }')"
ami_name="${name}-${architecture}-${input_hash:0:12}"

tofu init
tofu apply \
  -var "aws_region=$aws_region" \
  -var "name=$name" \
  -var "system=$system" \
  -var "builder_instance_type=$builder_instance_type" \
  -var "input_hash=$input_hash" \
  -var "builder_git_rev=$builder_git_rev" \
  -var "flake_lock_rev=$flake_lock_rev" \
  "${extra_tofu_args[@]}"

builder_instance_id="$(tofu output -raw builder_instance_id)"
nix_volume_id="$(tofu output -raw nix_volume_id)"
nix_device_name="$(tofu output -raw nix_device_name)"
nix_pool="$(tofu output -raw nix_pool)"

echo "Waiting for builder instance to stop: $builder_instance_id" >&2
aws ec2 wait instance-stopped --region "$aws_region" --instance-ids "$builder_instance_id"

capture_image_id="$(
  aws ec2 create-image \
    --region "$aws_region" \
    --instance-id "$builder_instance_id" \
    --name "${ami_name}-capture" \
    --description "Temporary capture for ${ami_name}" \
    --no-reboot \
    --tag-specifications "ResourceType=image,Tags=[{Key=Name,Value=${ami_name}-capture},{Key=Project,Value=${name}},{Key=Role,Value=ami-capture},{Key=InputHash,Value=${input_hash}}]" \
    --query ImageId \
    --output text
)"

aws ec2 wait image-available --region "$aws_region" --image-ids "$capture_image_id"

capture_image_json="$(aws ec2 describe-images --region "$aws_region" --image-ids "$capture_image_id")"
root_snapshot_id="$(
  jq -r --arg nixDevice "$nix_device_name" '
    .Images[0].BlockDeviceMappings[]
    | select(.DeviceName != $nixDevice and .Ebs != null)
    | .Ebs.SnapshotId
    ' <<<"$capture_image_json" | head -n 1
)"
root_device_name="$(
  jq -r --arg nixDevice "$nix_device_name" '
    .Images[0].BlockDeviceMappings[]
    | select(.DeviceName != $nixDevice and .Ebs != null)
    | .DeviceName
  ' <<<"$capture_image_json" | head -n 1
)"
root_volume_size="$(
  jq -r --arg device "$root_device_name" '
    .Images[0].BlockDeviceMappings[]
    | select(.DeviceName == $device)
    | .Ebs.VolumeSize
  ' <<<"$capture_image_json"
)"
nix_snapshot_id="$(
  jq -r --arg device "$nix_device_name" '
    .Images[0].BlockDeviceMappings[]
    | select(.DeviceName == $device)
    | .Ebs.SnapshotId
  ' <<<"$capture_image_json"
)"
nix_volume_size="$(
  jq -r --arg device "$nix_device_name" '
    .Images[0].BlockDeviceMappings[]
    | select(.DeviceName == $device)
    | .Ebs.VolumeSize
  ' <<<"$capture_image_json"
)"

if [ -z "$root_snapshot_id" ] || [ "$root_snapshot_id" = "null" ]; then
  echo "Unable to determine root snapshot for $capture_image_id" >&2
  exit 1
fi

if [ -z "$nix_snapshot_id" ] || [ "$nix_snapshot_id" = "null" ]; then
  echo "Unable to determine /nix snapshot for $capture_image_id" >&2
  exit 1
fi

block_device_mappings="$(
  jq -n \
    --arg rootDevice "$root_device_name" \
    --arg rootSnapshot "$root_snapshot_id" \
    --argjson rootSize "$root_volume_size" \
    --arg nixDevice "$nix_device_name" \
    --arg nixSnapshot "$nix_snapshot_id" \
    --argjson nixSize "$nix_volume_size" \
    '[
      {
        DeviceName: $rootDevice,
        Ebs: {
          SnapshotId: $rootSnapshot,
          VolumeSize: $rootSize,
          VolumeType: "gp3",
          DeleteOnTermination: true
        }
      },
      {
        DeviceName: $nixDevice,
        Ebs: {
          SnapshotId: $nixSnapshot,
          VolumeSize: $nixSize,
          VolumeType: "gp3",
          DeleteOnTermination: true
        }
      }
    ]'
)"

image_id="$(
  aws ec2 register-image \
    --region "$aws_region" \
    --name "$ami_name" \
    --description "NixOS EC2 AMI with matched root and ZFS /nix snapshots" \
    --architecture "$architecture" \
    --virtualization-type hvm \
    --root-device-name "$root_device_name" \
    --ena-support \
    --block-device-mappings "$block_device_mappings" \
    --query ImageId \
    --output text
)"

aws ec2 wait image-available --region "$aws_region" --image-ids "$image_id"

aws ec2 create-tags \
  --region "$aws_region" \
  --resources "$root_snapshot_id" \
  --tags \
    "Key=Name,Value=${ami_name}-root" \
    "Key=Project,Value=${name}" \
    "Key=Architecture,Value=${architecture}" \
    "Key=System,Value=${system}" \
    "Key=InputHash,Value=${input_hash}" \
    "Key=BuilderGitRev,Value=${builder_git_rev}" \
    "Key=FlakeLockRev,Value=${flake_lock_rev}"

aws ec2 create-tags \
  --region "$aws_region" \
  --resources "$nix_snapshot_id" "$nix_volume_id" \
  --tags \
    "Key=Name,Value=${ami_name}-nix" \
    "Key=Project,Value=${name}" \
    "Key=Architecture,Value=${architecture}" \
    "Key=System,Value=${system}" \
    "Key=InputHash,Value=${input_hash}" \
    "Key=BuilderGitRev,Value=${builder_git_rev}" \
    "Key=FlakeLockRev,Value=${flake_lock_rev}" \
    "Key=NixPool,Value=${nix_pool}"

aws ec2 create-tags \
  --region "$aws_region" \
  --resources "$image_id" \
  --tags \
    "Key=Name,Value=${ami_name}" \
    "Key=Project,Value=${name}" \
    "Key=Role,Value=published-ami" \
    "Key=Architecture,Value=${architecture}" \
    "Key=System,Value=${system}" \
    "Key=InputHash,Value=${input_hash}" \
    "Key=BuilderGitRev,Value=${builder_git_rev}" \
    "Key=FlakeLockRev,Value=${flake_lock_rev}" \
    "Key=NixPool,Value=${nix_pool}" \
    "Key=RootSnapshotId,Value=${root_snapshot_id}" \
    "Key=NixSnapshotId,Value=${nix_snapshot_id}"

aws ec2 deregister-image --region "$aws_region" --image-id "$capture_image_id" >/dev/null

jq -n \
  --arg architecture "$architecture" \
  --arg system "$system" \
  --arg inputHash "$input_hash" \
  --arg builderGitRev "$builder_git_rev" \
  --arg flakeLockRev "$flake_lock_rev" \
  --arg amiId "$image_id" \
  --arg rootSnapshotId "$root_snapshot_id" \
  --arg nixSnapshotId "$nix_snapshot_id" \
  --arg rootDeviceName "$root_device_name" \
  --arg nixDeviceName "$nix_device_name" \
  --arg nixPool "$nix_pool" \
  '{
    architecture: $architecture,
    system: $system,
    inputHash: $inputHash,
    builderGitRev: $builderGitRev,
    flakeLockRev: $flakeLockRev,
    amiId: $amiId,
    rootSnapshotId: $rootSnapshotId,
    nixSnapshotId: $nixSnapshotId,
    rootDeviceName: $rootDeviceName,
    nixDeviceName: $nixDeviceName,
    nixPool: $nixPool
  }' > artifact.json

if [ "$keep_builder" = false ]; then
  tofu destroy \
    -var "aws_region=$aws_region" \
    -var "name=$name" \
    -var "system=$system" \
    -var "builder_instance_type=$builder_instance_type" \
    -var "input_hash=$input_hash" \
    -var "builder_git_rev=$builder_git_rev" \
    -var "flake_lock_rev=$flake_lock_rev" \
    "${extra_tofu_args[@]}"
fi

echo "Published AMI: $image_id"