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

nix_device='${nix_device}'
nix_pool='${nix_pool}'
state_dir=/var/lib/nixos-zfs-ec2-ami
bootstrap_marker="$state_dir/bootstrap-generation-installed"
ready_marker="$state_dir/ami-ready"

mkdir -p "$state_dir"

write_bootstrap_config() {
  mkdir -p /etc/nixos
  cat > /etc/nixos/flake.nix <<'EOF'
{
  inputs = {
    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
  };

  outputs = { nixpkgs, ... }: {
    nixosConfigurations.ami-builder-bootstrap = nixpkgs.lib.nixosSystem {
      system = "${system}";
      modules = [
        ({ config, lib, modulesPath, pkgs, ... }: {
          imports = [ "$${modulesPath}/virtualisation/amazon-image.nix" ];

          networking.hostName = "nixos-zfs-ec2-builder";
          networking.hostId = "6e69787a";

          nix.settings.experimental-features = [ "nix-command" "flakes" ];
          boot.supportedFilesystems = [ "zfs" ];
          boot.zfs.forceImportRoot = false;

          environment.systemPackages = [
            pkgs.git
            pkgs.rsync
            pkgs.zfs
          ];

          services.amazon-ssm-agent.enable = true;

          users.users.builder-admin = {
            isNormalUser = true;
            extraGroups = [ "wheel" ];
            openssh.authorizedKeys.keys = ${ssh_keys};
          };
          security.sudo.wheelNeedsPassword = false;

          system.stateVersion = "25.05";
        })
      ];
    };
  };
}
EOF
}

write_final_config() {
  mkdir -p /etc/nixos
  cat > /etc/nixos/flake.nix <<'EOF'
{
  inputs = {
    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
  };

  outputs = { nixpkgs, ... }: {
    nixosConfigurations.nixos-zfs-ec2 = nixpkgs.lib.nixosSystem {
      system = "${system}";
      modules = [
        ({ config, lib, modulesPath, pkgs, ... }: {
          imports = [ "$${modulesPath}/virtualisation/amazon-image.nix" ];

          networking.hostName = "nixos-zfs-ec2";
          networking.hostId = "6e69787a";

          nix.settings.experimental-features = [ "nix-command" "flakes" ];
          boot.supportedFilesystems = [ "zfs" ];
          boot.zfs.extraPools = [ "${nix_pool}" ];
          boot.zfs.forceImportRoot = false;

          fileSystems."/nix" = {
            device = "${nix_pool}/nix";
            fsType = "zfs";
            neededForBoot = true;
          };

          systemd.sockets.nix-daemon = {
            after = [ "nix.mount" ];
            requires = [ "nix.mount" ];
          };

          systemd.services.nix-daemon = {
            after = [ "nix.mount" ];
            requires = [ "nix.mount" ];
          };

          environment.systemPackages = [
            pkgs.git
            pkgs.rsync
            pkgs.zfs
          ];

          services.amazon-ssm-agent.enable = true;
          systemd.services.amazon-init.enable = lib.mkForce false;

          systemd.services.nixos-user-data-flake = {
            description = "Apply NixOS flake from EC2 user data";
            wantedBy = [ "multi-user.target" ];
            after = [ "fetch-ec2-metadata.service" "network-online.target" "nix-daemon.socket" ];
            wants = [ "fetch-ec2-metadata.service" "network-online.target" "nix-daemon.socket" ];
            path = [ pkgs.coreutils pkgs.git ];
            serviceConfig = {
              Type = "oneshot";
              RemainAfterExit = true;
            };
            script = ''
              set -eu

              user_data=/etc/ec2-metadata/user-data
              state_dir=/var/lib/nixos-user-data-flake
              applied_hash_file="$state_dir/applied-sha256"

              if [ ! -s "$user_data" ]; then
                echo "No EC2 user-data flake found at $user_data" >&2
                exit 1
              fi

              mkdir -p "$state_dir" /etc/nixos
              new_hash="$(sha256sum "$user_data" | cut -d ' ' -f1)"
              old_hash="$(cat "$applied_hash_file" 2>/dev/null || true)"

              if [ "$new_hash" = "$old_hash" ]; then
                exit 0
              fi

              install -m 0644 "$user_data" /etc/nixos/flake.nix
              ${config.system.build.nixos-rebuild}/bin/nixos-rebuild switch --flake /etc/nixos
              printf '%s\n' "$new_hash" > "$applied_hash_file"
            '';
          };

          users.users.builder-admin = {
            isNormalUser = true;
            extraGroups = [ "wheel" ];
            openssh.authorizedKeys.keys = ${ssh_keys};
          };
          security.sudo.wheelNeedsPassword = false;

          system.stateVersion = "25.05";
        })
      ];
    };
  };
}
EOF
}

if [ ! -e "$bootstrap_marker" ]; then
  write_bootstrap_config
  nixos-rebuild boot --flake /etc/nixos#ami-builder-bootstrap
  touch "$bootstrap_marker"
  systemctl reboot
  exit 0
fi

for _ in $(seq 1 600); do
  if [ -e "$nix_device" ]; then
    break
  fi
  sleep 1
done

if [ ! -e "$nix_device" ]; then
  echo "Nix device did not appear: $nix_device" >&2
  exit 1
fi

if ! zpool list -H "$nix_pool" >/dev/null 2>&1; then
  if ! zpool import "$nix_pool" >/dev/null 2>&1; then
    zpool create \
      -f \
      -o ashift=12 \
      -O acltype=posixacl \
      -O atime=off \
      -O compression=zstd \
      -O mountpoint=none \
      -O xattr=sa \
      "$nix_pool" "$nix_device"

    zfs create -o mountpoint=legacy "$nix_pool/nix"
  fi
fi

mkdir -p /mnt/nixos-zfs-nix
if ! mountpoint -q /mnt/nixos-zfs-nix; then
  mount -t zfs "$nix_pool/nix" /mnt/nixos-zfs-nix
fi

write_final_config
nixos-rebuild boot --flake /etc/nixos#nixos-zfs-ec2

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 "$nix_pool"

cat > "$ready_marker" <<EOF
architecture=${system}
input_hash=${input_hash}
builder_git_rev=${builder_git_rev}
flake_lock_rev=${flake_lock_rev}
EOF

systemctl poweroff