{ config, lib, ... }:

{
  options.nanuqsaurus.impermanence.enable = lib.mkOption {
    type = lib.types.bool;
    default = false;
    description = "Enable impermanence with a /persist Btrfs subvolume.";
  };

  config = lib.mkIf config.nanuqsaurus.impermanence.enable {
    boot.initrd.supportedFilesystems = [ "btrfs" ];

    # Reset the /root subvolume to a pristine snapshot on every boot.
    # systemd stage-1 initrd does not support boot.initrd.postDeviceCommands, so
    # this runs as an initrd systemd service, ordered after the labelled device
    # appears and before the root filesystem is mounted (sysroot.mount).
    # Note: `cut` (coreutils) is used instead of `awk`, which is not in the
    # systemd initrd.
    boot.initrd.systemd.services.rollback = {
      description = "Rollback btrfs root subvolume to a pristine state";
      wantedBy = [ "initrd.target" ];
      after = [ "dev-disk-by\\x2dlabel-nixos.device" ];
      before = [ "sysroot.mount" ];
      unitConfig.DefaultDependencies = "no";
      serviceConfig.Type = "oneshot";
      script = ''
        mkdir -p /mnt
        mount -t btrfs -o subvol=/ /dev/disk/by-label/nixos /mnt

        if [ ! -d /mnt/root-blank ]; then
          btrfs subvolume snapshot -r /mnt/root /mnt/root-blank
        fi

        if [ -d /mnt/root ]; then
          if [ -d /mnt/root-previous ]; then
            btrfs subvolume delete /mnt/root-previous
          fi
          btrfs subvolume snapshot /mnt/root /mnt/root-previous

          while true; do
            subvols="$(btrfs subvolume list -o /mnt/root | cut -f9 -d' ')"
            if [ -z "$subvols" ]; then
              break
            fi
            for subvol in $subvols; do
              btrfs subvolume delete "/mnt/$subvol"
            done
          done

          btrfs subvolume delete /mnt/root
        fi

        btrfs subvolume snapshot /mnt/root-blank /mnt/root
        umount /mnt
      '';
    };

    fileSystems."/persist".neededForBoot = true;

    environment.persistence."/persist" = {
      hideMounts = true;
      directories = [
        "/etc/nixos"
        "/etc/NetworkManager"
        "/etc/ssh"
        "/var/lib/bluetooth"
        "/var/lib/cups"
        "/var/lib/fprint"
        "/var/lib/libvirt"
        "/var/lib/NetworkManager"
        "/var/lib/secureboot"
        "/var/lib/tailscale"
        "/var/lib/systemd"
        "/var/lib/nixos"
      ];
      files = [
        "/etc/machine-id"
      ];
    };
  };
}