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

aws_region="${AWS_REGION:-us-west-2}"
architecture="arm64"
bucket=""
subnet_id=""
security_group_id=""
instance_profile_name=""
instance_type="t4g.small"
nix_volume_size_gb="50"
name="nixos-zfs-ec2"

usage() {
  cat >&2 <<'EOF'
Usage: publish-ami.sh --bucket s3://bucket[/prefix] [options]

Options:
  --architecture <arm64|x86_64>       AMI architecture (default: arm64)
  --aws-region <region>               AWS region (default: AWS_REGION or us-west-2)
  --bucket <s3-uri>                   S3 URI used by EC2 VM import
  --subnet-id <subnet>                Subnet for the temporary seeder instance
  --security-group-id <sg>            Security group for the temporary seeder instance
  --instance-profile-name <profile>   Optional IAM instance profile for the seeder
  --instance-type <type>              Seeder instance type (default: t4g.small)
  --nix-volume-size-gb <gb>           /nix volume size (default: 50)
  --name <name>                       AMI name prefix
EOF
}

while [ "$#" -gt 0 ]; do
  case "$1" in
    --architecture)
      architecture="$2"
      shift 2
      ;;
    --aws-region)
      aws_region="$2"
      shift 2
      ;;
    --bucket)
      bucket="$2"
      shift 2
      ;;
    --subnet-id)
      subnet_id="$2"
      shift 2
      ;;
    --security-group-id)
      security_group_id="$2"
      shift 2
      ;;
    --instance-profile-name)
      instance_profile_name="$2"
      shift 2
      ;;
    --instance-type)
      instance_type="$2"
      shift 2
      ;;
    --nix-volume-size-gb)
      nix_volume_size_gb="$2"
      shift 2
      ;;
    --name)
      name="$2"
      shift 2
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "Unknown argument: $1" >&2
      usage
      exit 1
      ;;
  esac
done

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

if [ -z "$instance_profile_name" ]; then
  echo "Missing required --instance-profile-name for SSM seeding" >&2
  usage
  exit 1
fi

case "$architecture" in
  arm64)
    nix_system="aarch64-linux"
    ;;
  x86_64)
    nix_system="x86_64-linux"
    ;;
  *)
    echo "Unsupported architecture: $architecture" >&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 nix
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" "$architecture" | sha256sum | awk '{ print $1 }')"

root_image="$(nix build --print-out-paths ".#packages.${nix_system}.rootImage")"
nixos_system_closure="$(nix build --print-out-paths ".#packages.${nix_system}.systemClosure")"

tmpdir="$(mktemp -d)"
cleanup() {
  rm -rf "$tmpdir"
}
trap cleanup EXIT

bucket_no_scheme="${bucket#s3://}"
bucket_name="${bucket_no_scheme%%/*}"
bucket_prefix="${bucket_no_scheme#"$bucket_name"}"
bucket_prefix="${bucket_prefix#/}"
s3_key_prefix="${bucket_prefix:+$bucket_prefix/}${name}/${architecture}/${input_hash}"

image_file="$(find "$root_image" -type f \( -name '*.vhd' -o -name '*.raw' \) | head -n 1)"
if [ -z "$image_file" ]; then
  echo "Could not locate built root image under $root_image" >&2
  exit 1
fi

case "$image_file" in
  *.vhd)
    import_format="vhd"
    ;;
  *.raw)
    import_format="raw"
    ;;
  *)
    echo "Unsupported root image format: $image_file" >&2
    exit 1
    ;;
esac

aws s3 cp "$image_file" "s3://${bucket_name}/${s3_key_prefix}/root-image"

container_json="$tmpdir/container.json"
jq -n \
  --arg bucket "$bucket_name" \
  --arg key "${s3_key_prefix}/root-image" \
  --arg format "$import_format" \
  '[{Description:"nixos-zfs-ec2-root",Format:$format,UserBucket:{S3Bucket:$bucket,S3Key:$key}}]' \
  > "$container_json"

import_task_id="$(
  aws ec2 import-image \
    --region "$aws_region" \
    --architecture "$architecture" \
    --platform Linux \
    --disk-containers "file://$container_json" \
    --query ImportTaskId \
    --output text
)"

while true; do
  import_json="$(aws ec2 describe-import-image-tasks --region "$aws_region" --import-task-ids "$import_task_id")"
  status="$(jq -r '.ImportImageTasks[0].Status' <<<"$import_json")"
  case "$status" in
    completed)
      root_ami_id="$(jq -r '.ImportImageTasks[0].ImageId' <<<"$import_json")"
      break
      ;;
    deleted|deleting)
      echo "$import_json" >&2
      exit 1
      ;;
    *)
      status_message="$(jq -r '.ImportImageTasks[0].StatusMessage // ""' <<<"$import_json")"
      echo "import-image: $status $status_message" >&2
      sleep 20
      ;;
  esac
done

root_snapshot_id="$(
  aws ec2 describe-images \
    --region "$aws_region" \
    --image-ids "$root_ami_id" \
    --query 'Images[0].BlockDeviceMappings[?Ebs != `null`][0].Ebs.SnapshotId' \
    --output text
)"

run_args=(
  --region "$aws_region"
  --image-id "$root_ami_id"
  --instance-type "$instance_type"
  --block-device-mappings "DeviceName=/dev/sdf,Ebs={VolumeSize=${nix_volume_size_gb},VolumeType=gp3,DeleteOnTermination=true}"
  --tag-specifications "ResourceType=instance,Tags=[{Key=Name,Value=${name}-seed-${input_hash:0:12}}]"
  --query 'Instances[0].InstanceId'
  --output text
)

if [ -n "$subnet_id" ]; then
  run_args+=(--subnet-id "$subnet_id")
fi
if [ -n "$security_group_id" ]; then
  run_args+=(--security-group-ids "$security_group_id")
fi
if [ -n "$instance_profile_name" ]; then
  run_args+=(--iam-instance-profile "Name=$instance_profile_name")
fi

instance_id="$(aws ec2 run-instances "${run_args[@]}")"
aws ec2 wait instance-running --region "$aws_region" --instance-ids "$instance_id"
aws ec2 wait instance-status-ok --region "$aws_region" --instance-ids "$instance_id"

cat > "$tmpdir/seed-nix.sh" <<'EOS'
#!/usr/bin/env bash
set -euo pipefail

pool="nixos-zfs-nix"
device=""
for candidate in /dev/disk/by-id/nvme-Amazon_Elastic_Block_Store_*; do
  if [ -e "$candidate" ] && ! findmnt -n -S "$candidate" >/dev/null 2>&1; then
    device="$candidate"
  fi
done

if [ -z "$device" ]; then
  echo "Unable to find unattached EBS device for /nix seed" >&2
  exit 1
fi

zpool create \
  -f \
  -o ashift=12 \
  -O acltype=posixacl \
  -O atime=off \
  -O compression=zstd \
  -O mountpoint=none \
  -O xattr=sa \
  "$pool" "$device"

zfs create -o mountpoint=legacy "$pool/nix"
mkdir -p /mnt/nixos-zfs-nix
mount -t zfs "$pool/nix" /mnt/nixos-zfs-nix
rsync -aHAX --numeric-ids \
  --exclude='/var/nix/daemon-socket/*' \
  /nix/ /mnt/nixos-zfs-nix/
rm -f /mnt/nixos-zfs-nix/var/nix/daemon-socket/socket
umount /mnt/nixos-zfs-nix
zpool export "$pool"
EOS

jq -n --rawfile commands "$tmpdir/seed-nix.sh" '{commands: [$commands]}' > "$tmpdir/ssm-parameters.json"

command_id="$(
  aws ssm send-command \
    --region "$aws_region" \
    --instance-ids "$instance_id" \
    --document-name AWS-RunShellScript \
    --parameters "file://$tmpdir/ssm-parameters.json" \
    --query Command.CommandId \
    --output text
)"

while true; do
  invocation_json="$(aws ssm get-command-invocation --region "$aws_region" --command-id "$command_id" --instance-id "$instance_id" 2>/dev/null || true)"
  status="$(jq -r '.Status // "Pending"' <<<"$invocation_json")"
  case "$status" in
    Success)
      break
      ;;
    Failed|Cancelled|TimedOut)
      echo "$invocation_json" >&2
      exit 1
      ;;
    *)
      echo "seed /nix: $status" >&2
      sleep 10
      ;;
  esac
done

aws ec2 stop-instances --region "$aws_region" --instance-ids "$instance_id" >/dev/null
aws ec2 wait instance-stopped --region "$aws_region" --instance-ids "$instance_id"

nix_volume_id="$(
  aws ec2 describe-instances \
    --region "$aws_region" \
    --instance-ids "$instance_id" \
    --query 'Reservations[0].Instances[0].BlockDeviceMappings[?DeviceName==`/dev/sdf`].Ebs.VolumeId | [0]' \
    --output text
)"

nix_snapshot_id="$(
  aws ec2 create-snapshot \
    --region "$aws_region" \
    --volume-id "$nix_volume_id" \
    --description "${name} ${architecture} ZFS /nix ${input_hash}" \
    --tag-specifications "ResourceType=snapshot,Tags=[{Key=Name,Value=${name}-nix-${architecture}-${input_hash:0:12}},{Key=InputHash,Value=${input_hash}},{Key=BuilderGitRev,Value=${builder_git_rev}}]" \
    --query SnapshotId \
    --output text
)"
aws ec2 wait snapshot-completed --region "$aws_region" --snapshot-ids "$nix_snapshot_id"

tofu init
tofu apply \
  -var "aws_region=$aws_region" \
  -var "name=$name" \
  -var "architecture=$architecture" \
  -var "root_snapshot_id=$root_snapshot_id" \
  -var "nix_snapshot_id=$nix_snapshot_id" \
  -var "input_hash=$input_hash" \
  -var "builder_git_rev=$builder_git_rev" \
  -var "flake_lock_rev=$flake_lock_rev" \
  -var "nixos_system_closure=$nixos_system_closure" \
  -var "nix_volume_size_gb=$nix_volume_size_gb"

ami_id="$(tofu output -raw ami_id)"

aws ec2 terminate-instances --region "$aws_region" --instance-ids "$instance_id" >/dev/null
aws ec2 deregister-image --region "$aws_region" --image-id "$root_ami_id" >/dev/null || true

jq -n \
  --arg architecture "$architecture" \
  --arg inputHash "$input_hash" \
  --arg builderGitRev "$builder_git_rev" \
  --arg flakeLockRev "$flake_lock_rev" \
  --arg amiId "$ami_id" \
  --arg rootSnapshotId "$root_snapshot_id" \
  --arg nixSnapshotId "$nix_snapshot_id" \
  --arg nixosSystemClosure "$nixos_system_closure" \
  '{
    architecture: $architecture,
    inputHash: $inputHash,
    builderGitRev: $builderGitRev,
    flakeLockRev: $flakeLockRev,
    amiId: $amiId,
    rootSnapshotId: $rootSnapshotId,
    nixSnapshotId: $nixSnapshotId,
    nixosSystemClosure: $nixosSystemClosure
  }' > artifact.json

echo "Published AMI: $ami_id"