nix-config/modules/nixos-modules/storage/storage.nix

args @ {
  lib,
  config,
  ...
}: let
  datasetSubmodule = (import ./submodules/dataset.nix) args;
  impermanenceDatasetSubmodule = (import ./submodules/impermanenceDataset.nix) args;

  # Get the option names from both submodules to automatically determine which are impermanence-specific
  regularDatasetEval = lib.evalModules {
    modules = [datasetSubmodule];
    specialArgs = args;
  };
  impermanenceDatasetEval = lib.evalModules {
    modules = [impermanenceDatasetSubmodule];
    specialArgs = args;
  };

  regularDatasetOptions = builtins.attrNames regularDatasetEval.options;
  impermanenceDatasetOptions = builtins.attrNames impermanenceDatasetEval.options;

  # Find options that are only in impermanence datasets (not in regular ZFS datasets)
  impermanenceOnlyOptions = lib.lists.subtractLists regularDatasetOptions impermanenceDatasetOptions;
in {
  options.storage.datasets = {
    ephemeral = lib.mkOption {
      type = lib.types.attrsOf (lib.types.submodule datasetSubmodule);
      default = {};
    };
    local = lib.mkOption {
      type = lib.types.attrsOf (lib.types.submodule impermanenceDatasetSubmodule);
      default = {};
    };
    replicate = lib.mkOption {
      type = lib.types.attrsOf (lib.types.submodule impermanenceDatasetSubmodule);
      default = {};
    };
  };

  config = lib.mkMerge [
    (lib.mkIf config.storage.zfs.enable {
      # Create ZFS datasets based on storage.datasets configuration
      storage.datasets = {
        local = {
          "nix" = {
            type = "zfs_fs";
            mount = "/nix";
            snapshot = {
              autoSnapshot = false;
            };
            atime = "off";
            relatime = "off";
          };
        };
        replicate = {
          "system/var/log" = {
            type = "zfs_fs";
            mount = "/var/log";
          };
        };
      };
    })
    (lib.mkIf (config.storage.zfs.enable && config.storage.impermanence.enable) {
      storage.datasets = {
        ephemeral = {
          "" = {
            type = "zfs_fs";
            mount = null;
          };
          "system/root" = {
            type = "zfs_fs";
            mount = "/";
            snapshot = {
              blankSnapshot = true;
            };
          };
        };
        # TODO: can we auto set the mount points on these to just be `"/persist/local/${name}"`
        local = {
          "" = {
            mount = "/persist/local/";
          };
        };
        # TODO: can we auto set the mount points on these to just be `"/persist/replicate/${name}"`
        replicate = {
          "" = {
            mount = "/persist/replicate/";
          };
          "system/root" = {
            mount = "/persist/replicate/system/root";
            snapshot = {
              autoSnapshot = true;
            };
            directories = {
              "/var/lib/nixos".enable = true;
              "/var/lib/systemd/coredump".enable = true;
            };
            files = {
              "/etc/machine-id".enable = true;
            };
          };
          "home" = {
            mount = "/persist/replicate/home";
            snapshot = {
              autoSnapshot = true;
            };
          };
        };
      };

      storage.zfs.datasets = lib.mkMerge [
        (lib.mapAttrs' (name: dataset: {
            name =
              if name == ""
              then "ephemeral"
              else "ephemeral/${name}";
            value = dataset;
          })
          config.storage.datasets.ephemeral)
      ];

      boot.initrd.postResumeCommands = lib.mkAfter ''
        zfs rollback -r rpool/ephemeral/system/root@blank
      '';

      storage.impermanence.datasets = lib.mkMerge [
        (lib.mapAttrs' (name: dataset: {
            name =
              if name == ""
              then "persist/local"
              else "persist/local/${name}";
            value = dataset;
          })
          config.storage.datasets.local)
        (lib.mapAttrs' (name: dataset: {
            name =
              if name == ""
              then "persist/replicate"
              else "persist/replicate/${name}";
            value = dataset;
          })
          config.storage.datasets.replicate)
      ];
    })
    (lib.mkIf (config.storage.zfs.enable && !config.storage.impermanence.enable) {
      storage.datasets = {
        # Base organizational datasets (only needed when impermanence is disabled)
        local = {
          "" = {
            type = "zfs_fs";
            mount = "";
          };
          "root" = {
            type = "zfs_fs";
            mount = "/";
            compression = "lz4";
            acltype = "posixacl";
            relatime = "on";
            xattr = "sa";
            snapshot = {
              autoSnapshot = true;
              blankSnapshot = true;
            };
          };
        };
        replicate = {
          "" = {
            type = "zfs_fs";
            mount = null;
          };
          "system/var/log" = {
            type = "zfs_fs";
            mount = "/var/log";
          };
        };
      };

      storage.zfs.datasets = lib.mkMerge [
        (lib.mapAttrs' (name: dataset: {
            name =
              if name == ""
              then "persist/local"
              else "persist/local/${name}";
            value = builtins.removeAttrs dataset impermanenceOnlyOptions;
          })
          config.storage.datasets.local)
        (lib.mapAttrs' (name: dataset: {
            name =
              if name == ""
              then "persist/replicate"
              else "persist/replicate/${name}";
            value = builtins.removeAttrs dataset impermanenceOnlyOptions;
          })
          config.storage.datasets.replicate)
      ];
    })
  ];

  # TODO: set up datasets for systemd services that want a dataset created
  # TODO: home-manager.users.<user>.storage.impermanence.enable
  #   is false then persist the entire directory of the user
  #   if true persist home-manager.users.<user>.storage.impermanence.datasets
  # TODO: systemd.services.<name>.storage.datasets persists
  # TODO: configure other needed storage modes here
}