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

  # Hash function for disk names (max 27 chars to fit GPT limitations)
  hashDisk = drive: (builtins.substring 0 27 (builtins.hashString "sha256" drive));

  # Helper to flatten vdevs into list of devices with names
  allVdevDevices = lib.lists.flatten (builtins.map (
      vdev:
        builtins.map (
          device:
            lib.attrsets.nameValuePair (hashDisk device.device) device
        )
        vdev
    )
    config.storage.zfs.pool.vdevs);

  # Cache devices with names
  allCacheDevices = builtins.map (
    device:
      lib.attrsets.nameValuePair (hashDisk device.device) device
  ) (config.storage.zfs.pool.cache);

  # All devices (vdevs + cache)
  allDevices = allVdevDevices ++ allCacheDevices;

  # Boot devices - filter devices that have boot = true
  bootDevices = builtins.filter (device: device.value.boot) allDevices;

  # Helper function to convert dataset options to ZFS properties
  datasetToZfsOptions = dataset: let
    baseOptions =
      (lib.attrsets.optionalAttrs (dataset.acltype != null) {acltype = dataset.acltype;})
      // (lib.attrsets.optionalAttrs (dataset.relatime != null) {relatime = dataset.relatime;})
      // (lib.attrsets.optionalAttrs (dataset.atime != null) {atime = dataset.atime;})
      // (lib.attrsets.optionalAttrs (dataset.xattr != null) {xattr = dataset.xattr;})
      // (lib.attrsets.optionalAttrs (dataset.compression != null) {compression = dataset.compression;})
      // (lib.attrsets.optionalAttrs (dataset.sync != null) {sync = dataset.sync;})
      // (lib.attrsets.optionalAttrs (dataset.recordSize != null) {recordSize = dataset.recordSize;});

    encryptionOptions = lib.attrsets.optionalAttrs (dataset.encryption.enable) (
      (lib.attrsets.optionalAttrs (dataset.encryption ? type) {encryption = dataset.encryption.type;})
      // (lib.attrsets.optionalAttrs (dataset.encryption ? keyformat) {keyformat = dataset.encryption.keyformat;})
      // (lib.attrsets.optionalAttrs (dataset.encryption ? keylocation) {keylocation = dataset.encryption.keylocation;})
    );

    mountOptions = lib.attrsets.optionalAttrs (dataset ? mount && dataset.mount ? enable) (
      if builtins.isBool dataset.mount.enable
      then {
        canmount =
          if dataset.mount.enable
          then "on"
          else "off";
      }
      else {canmount = dataset.mount.enable;}
    );

    snapshotOptions = lib.attrsets.optionalAttrs (dataset ? snapshot && dataset.snapshot ? autoSnapshot) {
      "com.sun:auto-snapshot" =
        if dataset.snapshot.autoSnapshot
        then "true"
        else "false";
    };
  in
    baseOptions // encryptionOptions // mountOptions // snapshotOptions;

  # Helper to generate post create hooks
  generatePostCreateHook = name: dataset:
    dataset.postCreateHook
    + (lib.optionalString dataset.snapshot.blankSnapshot ''
      zfs snapshot rpool/${name}@blank
    '');

  # Convert datasets to disko format
  convertedDatasets = builtins.listToAttrs (
    (lib.attrsets.mapAttrsToList (
        name: dataset:
          lib.attrsets.nameValuePair name {
            type = dataset.type;
            options = datasetToZfsOptions dataset;
            mountpoint = dataset.mount.mountPoint or null;
            postCreateHook = generatePostCreateHook name dataset;
          }
      )
      config.storage.zfs.datasets)
    ++ (lib.optional (config.storage.zfs.rootDataset != null) (
      lib.attrsets.nameValuePair "" {
        type = config.storage.zfs.rootDataset.type;
        options = datasetToZfsOptions config.storage.zfs.rootDataset;
        mountpoint = config.storage.zfs.rootDataset.mount.mountPoint or null;
        postCreateHook = generatePostCreateHook "" config.storage.zfs.rootDataset;
      }
    ))
  );
in {
  options.storage = {
    zfs = {
      enable = lib.mkEnableOption "Should zfs be enabled on this system.";

      notifications = {
        enable = lib.mkEnableOption "are notifications enabled";
        host = lib.mkOption {
          type = lib.types.str;
          description = "what is the host that we are going to send the email to";
        };
        port = lib.mkOption {
          type = lib.types.port;
          description = "what port is the host using to receive mail on";
        };
        to = lib.mkOption {
          type = lib.types.str;
          description = "what account is the email going to be sent to";
        };
        user = lib.mkOption {
          type = lib.types.str;
          description = "what user is the email going to be set from";
        };
        tokenFile = lib.mkOption {
          type = lib.types.str;
          description = "file containing the password to be used by msmtp for notifications";
        };
      };

      pool = let
        deviceType =
          lib.types.coercedTo lib.types.str (device: {
            device = device;
            boot = false;
          }) (lib.types.submodule {
            options = {
              device = lib.mkOption {
                type = lib.types.str;
              };
              boot = lib.mkEnableOption "should this device be a boot device";
            };
          });
      in {
        encryption = {
          enable = lib.mkEnableOption "Should encryption be enabled on this pool.";
          keyformat = lib.mkOption {
            type = lib.types.enum ["raw" "hex" "passphrase"];
            default = "hex";
            description = "Format of the encryption key";
          };
          keylocation = lib.mkOption {
            type = lib.types.str;
            default = "prompt";
            description = "Location of the encryption key";
          };
        };
        mode = lib.mkOption {
          type = lib.types.enum ["stripe" "mirror" "raidz1" "raidz2" "raidz3"];
          default = "raidz2";
          description = "ZFS redundancy mode for the pool";
        };
        bootPartitionSize = lib.mkOption {
          type = lib.types.str;
          default = "2G";
          description = "Size of the boot partition on boot drives";
        };
        vdevs = lib.mkOption {
          type = lib.types.listOf (lib.types.listOf deviceType);
          default = [];
          description = "List of vdevs, where each vdev is a list of devices";
        };
        cache = lib.mkOption {
          type = lib.types.listOf deviceType;
          default = [];
        };
      };

      rootDataset = lib.mkOption {
        type = lib.types.nullOr (lib.types.submodule datasetSubmodule);
        description = "Root ZFS dataset to create";
        default = null;
      };

      datasets = lib.mkOption {
        type = lib.types.attrsOf (lib.types.submodule datasetSubmodule);
        description = "Additional ZFS datasets to create";
        default = {};
      };
    };
  };

  config = lib.mkIf config.storage.zfs.enable (lib.mkMerge [
    {
      # Assertion that we have at least one boot device
      assertions = [
        {
          assertion = (builtins.length bootDevices) > 0;
          message = "ZFS configuration requires at least one boot device. Set boot = true for at least one device in your vdevs or cache.";
        }
      ];

      # # Warning about disk/dataset mismatches - these would be runtime checks
      # warnings = let
      #   configuredDisks = builtins.map (device: device.device) (builtins.map (dev: dev.value) allDevices);
      #   diskWarnings =
      #     lib.optional (config.storage.zfs.enable)
      #     "ZFS: Please ensure the following disks are available on your system: ${builtins.concatStringsSep ", " configuredDisks}";

      #   configuredDatasets = builtins.attrNames config.storage.zfs.datasets;
      #   datasetWarnings =
      #     lib.optional (config.storage.zfs.enable && (builtins.length configuredDatasets) > 0)
      #     "ZFS: Configured datasets: ${builtins.concatStringsSep ", " configuredDatasets}. Ensure these match your intended ZFS layout.";
      # in
      #   diskWarnings ++ datasetWarnings;

      # services.zfs = {
      #   autoScrub.enable = true;
      #   autoSnapshot.enable = true;
      # };

      # # Configure disko for ZFS setup
      disko.devices = {
        disk = builtins.listToAttrs (
          builtins.map (
            drive:
              lib.attrsets.nameValuePair (drive.name) {
                type = "disk";
                device = "/dev/disk/by-id/${drive.value.device}";
                content = {
                  type = "gpt";
                  partitions = {
                    ESP = lib.mkIf drive.value.boot {
                      size = config.storage.zfs.pool.bootPartitionSize;
                      type = "EF00";
                      content = {
                        type = "filesystem";
                        format = "vfat";
                        mountpoint = "/boot";
                        mountOptions = ["umask=0077"];
                      };
                    };
                    zfs = {
                      size = "100%";
                      content = {
                        type = "zfs";
                        pool = "rpool";
                      };
                    };
                  };
                };
              }
          )
          allDevices
        );

        zpool = {
          rpool = {
            type = "zpool";
            mode = {
              topology = {
                type = "topology";
                vdev =
                  builtins.map (vdev: {
                    mode = config.storage.zfs.pool.mode;
                    members = builtins.map (device: hashDisk device.device) vdev;
                  })
                  config.storage.zfs.pool.vdevs;
                cache = builtins.map (device: hashDisk device.device) config.storage.zfs.pool.cache;
              };
            };

            options = {
              ashift = "12";
              autotrim = "on";
            };

            rootFsOptions =
              {
                canmount = "off";
                mountpoint = "none";
                xattr = "sa";
                acltype = "posixacl";
                relatime = "on";
                compression = "lz4";
                "com.sun:auto-snapshot" = "false";
              }
              // (lib.attrsets.optionalAttrs config.storage.zfs.pool.encryption.enable {
                encryption = "on";
                keyformat = config.storage.zfs.pool.encryption.keyformat;
                keylocation = config.storage.zfs.pool.encryption.keylocation;
              });

            datasets = convertedDatasets;
          };
        };
      };
    }
    (lib.mkIf config.storage.zfs.notifications.enable {
      programs.msmtp = {
        enable = true;
        setSendmail = true;
        defaults = {
          aliases = "/etc/aliases";
          port = config.storage.zfs.notifications.port;
          tls_trust_file = "/etc/ssl/certs/ca-certificates.crt";
          tls = "on";
          auth = "login";
          tls_starttls = "off";
        };
        accounts = {
          zfs_notifications = {
            auth = true;
            tls = true;
            host = config.storage.zfs.notifications.host;
            passwordeval = "cat ${config.storage.zfs.notifications.tokenFile}";
            user = config.storage.zfs.notifications.user;
            from = config.storage.zfs.notifications.user;
          };
        };
      };

      services.zfs = {
        zed = {
          enableMail = true;

          settings = {
            ZED_DEBUG_LOG = "/tmp/zed.debug.log";
            ZED_EMAIL_ADDR = [config.storage.zfs.notifications.to];
            ZED_EMAIL_PROG = "${pkgs.msmtp}/bin/msmtp";
            ZED_EMAIL_OPTS = "-a zfs_notifications @ADDRESS@";

            ZED_NOTIFY_INTERVAL_SECS = 3600;
            ZED_NOTIFY_VERBOSE = true;

            ZED_USE_ENCLOSURE_LEDS = true;
            ZED_SCRUB_AFTER_RESILVER = true;
          };
        };
      };
    })
  ]);
}