{ inputs, ... }:
{
  flake.modules.nixos.base =
    {
      config,
      lib,
      pkgs,
      ...
    }:
    {
      imports = [ inputs.impermanence.nixosModules.impermanence ];

      options.infra.persist = {
        dir = lib.mkOption {
          type = lib.types.path;
          default = "/persist";
          description = "The root directory to use for persistence.";
        };

        directories = lib.mkOption {
          type = lib.types.listOf lib.types.str;
          default = [ ];
          description = "The list of root owned directories to persist across reboots.";
        };

        files = lib.mkOption {
          type = lib.types.listOf lib.types.str;
          default = [ ];
          description = "The list of root owned files to persist across reboots.";
        };

        users = lib.mkOption {
          type = lib.types.attrsOf (
            lib.types.submodule {
              options = {
                directories = lib.mkOption {
                  type = lib.types.listOf lib.types.str;
                  default = [ ];
                  description = "The list of directories to persist within the users home directory.";
                };

                files = lib.mkOption {
                  type = lib.types.listOf lib.types.str;
                  default = [ ];
                  description = "The list of files to persist within the users home directory.";
                };
              };
            }
          );
          default = { };
          description = "The directories and files to persist for a user.";
        };
      };

      config =
        let
          cfg = config.infra.persist;
          blank = "${config.infra.zfs.pool}/local/root@blank";
        in
        {
          boot = {
            initrd = {
              supportedFilesystems.zfs = true;
              systemd = {
                storePaths = [ pkgs.zfs ];
                services.infra-rollback-root = {
                  description = "Rollback root ZFS dataset to blank snapshot";
                  wantedBy = [ "initrd.target" ];
                  after = [ "zfs-import-${config.infra.zfs.pool}.service" ];
                  before = [ "sysroot.mount" ];
                  enableStrictShellChecks = true;

                  serviceConfig = {
                    Type = "oneshot";
                  };

                  script = ''
                    if "${pkgs.zfs}/bin/zfs" list -H -o name "${blank}" >/dev/null 2>&1; then
                      "${pkgs.zfs}/bin/zfs" rollback -r "${blank}"
                    fi
                  '';
                };
              };
            };
          };

          fileSystems."${cfg.dir}".neededForBoot = true;
          environment.persistence."${cfg.dir}" = {
            directories = [
              # keep-sorted start
              "/var/lib/nixos"
              "/var/lib/systemd"
              "/var/log/journal"
              # keep-sorted end
            ]
            ++ cfg.directories;
            files = [
              "/etc/machine-id"

              # Needed for sops-nix to be able to decrypt secrets.
              "/etc/ssh/ssh_host_ed25519_key"
              "/etc/ssh/ssh_host_ed25519_key.pub"
            ]
            ++ cfg.files;
          };

          assertions = [
            {
              assertion = config.boot.initrd.systemd.enable;
              message = "initrd systemd must be enabled for persistence (impermanence)";
            }
          ];
        };
    };
}