{ config
, inputs
, pkgs
, ...
}:
{
  imports = [
    inputs.impermanence.nixosModules.impermanence
  ];

  boot.initrd.systemd.services.rollback-root = {
    description = "Rollback Btrfs root filesystem to a pristine state";
    wantedBy = [ "initrd.target" ];
    
    # Run after the physical storage device node exists
    after = [ "block-device-systemd-sysext.target" ]; 
    # Run strictly before the root partition gets mounted by stage 1
    before = [ "sysroot.mount" ]; 
    
    unitConfig.DefaultDependencies = "no";
    
    # We explicitly include packages needed inside the initrd environment
    path = with pkgs; [ btrfs-progs coreutils util-linux gawk gnugrep ];

    serviceConfig = {
      Type = "oneshot";
      RemainAfterExit = true;
      ExecStart = pkgs.writeShellScript "rollback-root-script" ''
        set -e
        DEVICE="${config.disko.devices.disk.main.device}-part2"

        # Mount Btrfs root
        mkdir -p /mnt
        if ! mount -o subvol=/ "$DEVICE" /mnt; then
          echo "Failed to mount $DEVICE at /mnt"
          exit 1
        fi

        # Create directory for old roots
        mkdir -p /mnt/old-roots

        # Move current root to old-roots with current timestamp
        if [[ -e /mnt/@root ]]; then
          timestamp=$(date +%Y-%m-%d_%H:%M:%S)
          if ! btrfs subvolume snapshot -r /mnt/@root "/mnt/old-roots/@root-$timestamp"; then
            echo "Failed to move /mnt/@root to old-roots"
            umount /mnt
            exit 1
          fi
        fi

        # Function to recursively delete subvolumes
        delete_subvolume_recursively() {
          for i in $(btrfs subvolume list -o "$1" | cut -f 9- -d ' '); do
            delete_subvolume_recursively "/mnt/$i"
          done
          if ! btrfs subvolume delete "$1"; then
            echo "Failed to delete subvolume $1"
          fi
        }

        # Delete old roots more than 5
        index=0
        for i in $(btrfs subvolume list /mnt | grep 'old-roots/@root-' | cut -f 9- -d ' ' | sort -r); do
          if [[ $index -ge 5 ]]; then
            delete_subvolume_recursively "/mnt/$i"
          fi
          index=$((index + 1))
        done

        # Create or restore fresh root
        if [[ -e /mnt/@root-blank ]]; then
          delete_subvolume_recursively /mnt/@root
          if ! btrfs subvolume snapshot /mnt/@root-blank /mnt/@root; then
            echo "Failed to snapshot @root-blank to @root"
            umount /mnt
            exit 1
          fi
        else
          if ! btrfs subvolume create /mnt/@root-blank; then
            echo "Failed to create @root-blank"
            umount /mnt
            exit 1
          fi
          if ! btrfs subvolume create /mnt/@root; then
            echo "Failed to create @root"
            umount /mnt
            exit 1
          fi
        fi

        # Unmount
        if ! umount /mnt; then
          echo "Failed to unmount /mnt"
          exit 1
        fi
      '';
    };
  };

  # Persistent directories for impermanence
  fileSystems."/persist".neededForBoot = true;
  fileSystems."/var/lib".neededForBoot = true;

  environment.persistence."/persist" = {
    hideMounts = true;

    directories = [
      "/etc/ssh"
      "/var/cache"
    ];

    files = [
      "/etc/machine-id"
    ];

    users.thomas.directories = [
      ".ssh"
      ".local/share"
    ];
  };
}