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

  # max gpt length is 36 and disk adds formats it like disk-xxxx-zfs which means we need to be 9 characters under that
  hashDisk = drive: (builtins.substring 0 27 (builtins.hashString "sha256" drive));

  poolVdevs =
    builtins.map (
      vdev:
        builtins.map (
          device: let
            deviceStr =
              if builtins.isString device
              then device
              else device.device;
          in
            lib.attrsets.nameValuePair (hashDisk deviceStr) deviceStr
        )
        vdev
    )
    config.storage.zfs.pool.vdevs;

  poolCache = builtins.map (
    name: let
      device = config.storage.zfs.pool.cache.${name};
      deviceStr =
        if builtins.isString device
        then device
        else device.device;
    in
      lib.attrsets.nameValuePair (hashDisk deviceStr) deviceStr
  ) (builtins.attrNames config.storage.zfs.pool.cache);

  bootDrives =
    builtins.map (
      device:
        if builtins.isString device
        then device
        else device.device
    ) (builtins.filter (
        device:
          if builtins.isString device
          then false
          else device.boot
      )
      (lib.lists.flatten config.storage.zfs.pool.vdevs));

  allDrives = (lib.lists.flatten poolVdevs) ++ poolCache;
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;
          }) {
            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.attrsOf deviceType;
          default = {};
        };
      };

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

      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 [
    {
      assertions = [
        {
          assertion = builtins.length bootDrives > 0;
          message = ''
            ZFS configuration requires at least one boot drive. Please configure at least one device with boot = true in storage.zfs.pool.vdevs.
          '';
        }
        {
          assertion =
            !(
              config.storage.zfs.pool.encryption.enable
              && (config.storage.zfs.rootDataset.encryption
                != null
                || config.storage.zfs.rootDataset.keyformat != null
                || config.storage.zfs.rootDataset.keylocation != null)
            );
          message = ''
            Cannot set encryption options in both pool.encryption and rootDataset.
            Use either pool.encryption for default settings or rootDataset encryption options for explicit control, but not both.
          '';
        }
      ];

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

      # Disko configuration based on pool settings
      # disko.devices = {
      #   disk = (
      #     builtins.listToAttrs (
      #       builtins.map
      #       (drive:
      #         lib.attrsets.nameValuePair (drive.name) {
      #           type = "disk";
      #           device = "/dev/disk/by-id/${drive.value}";
      #           content = {
      #             type = "gpt";
      #             partitions = {
      #               ESP = lib.mkIf (builtins.elem drive.value bootDrives) {
      #                 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";
      #                 };
      #               };
      #             };
      #           };
      #         })
      #       allDrives
      #     )
      #   );
      #   zpool = {
      #     rpool = {
      #       type = "zpool";
      #       mode = {
      #         topology = {
      #           type = "topology";
      #           vdev = (
      #             builtins.map (disks: {
      #               mode = config.storage.zfs.pool.mode;
      #               members =
      #                 builtins.map (disk: disk.name) disks;
      #             })
      #             poolVdevs
      #           );
      #           cache = builtins.map (disk: disk.name) poolCache;
      #         };
      #       };

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

      #       rootFsOptions = let
      #         rootDataset = config.storage.zfs.rootDataset;
      #         # Start with defaults that match the original hardcoded values
      #         defaults = {
      #           canmount = "off";
      #           mountpoint = "none";
      #           xattr = "sa";
      #           acltype = "posixacl";
      #           relatime = "on";
      #           compression = "lz4";
      #           "com.sun:auto-snapshot" = "false";
      #         };
      #         # Override defaults with non-null values from rootDataset
      #         userOptions = lib.attrsets.filterAttrs (_: v: v != null) {
      #           canmount = rootDataset.canmount;
      #           mountpoint = rootDataset.mountpoint;
      #           xattr = rootDataset.xattr;
      #           acltype = rootDataset.acltype;
      #           relatime = rootDataset.relatime;
      #           compression = rootDataset.compression;
      #           encryption = rootDataset.encryption;
      #           keyformat = rootDataset.keyformat;
      #           keylocation = rootDataset.keylocation;
      #           recordsize = rootDataset.recordsize;
      #           sync = rootDataset.sync;
      #           atime = rootDataset.atime;
      #           "com.sun:auto-snapshot" =
      #             if rootDataset.autoSnapshot == null
      #             then null
      #             else
      #               (
      #                 if rootDataset.autoSnapshot
      #                 then "true"
      #                 else "false"
      #               );
      #         };
      #         # Only apply pool encryption if user hasn't set encryption options in rootDataset
      #         poolEncryptionOptions =
      #           lib.attrsets.optionalAttrs (
      #             config.storage.zfs.pool.encryption.enable
      #             && rootDataset.encryption == null
      #             && rootDataset.keyformat == null
      #             && rootDataset.keylocation == null
      #           ) {
      #             encryption = "on";
      #             keyformat = config.storage.zfs.pool.encryption.keyformat;
      #             keylocation = config.storage.zfs.pool.encryption.keylocation;
      #           };
      #       in
      #         defaults // userOptions // rootDataset.options // poolEncryptionOptions;

      #       datasets = lib.mkMerge [
      #         (
      #           lib.attrsets.mapAttrs (name: value: {
      #             type = value.type;
      #             options = let
      #               # For datasets, only include non-null user-specified values
      #               userOptions = lib.attrsets.filterAttrs (_: v: v != null) {
      #                 canmount = value.canmount;
      #                 xattr = value.xattr;
      #                 acltype = value.acltype;
      #                 relatime = value.relatime;
      #                 compression = value.compression;
      #                 encryption = value.encryption;
      #                 keyformat = value.keyformat;
      #                 keylocation = value.keylocation;
      #                 recordsize = value.recordsize;
      #                 sync = value.sync;
      #                 atime = value.atime;
      #                 "com.sun:auto-snapshot" =
      #                   if value.autoSnapshot == null
      #                   then null
      #                   else
      #                     (
      #                       if value.autoSnapshot
      #                       then "true"
      #                       else "false"
      #                     );
      #               };
      #             in
      #               userOptions // (value.options or {});
      #             mountpoint = value.mountpoint;
      #             postCreateHook = value.postCreateHook or "";
      #           })
      #           config.storage.zfs.datasets
      #         )
      #       ];
      #     };
      #   };
      # };

      # Post-activation scripts for validation
      system.activationScripts = {
        # Script 1: Validate pool, cache devices, and vdevs
        zfs-pool-validation = {
          text = ''
            echo "Running ZFS pool validation..."

            # Function to check if a device exists in a vdev or cache
            check_device_in_pool() {
              local device_id="$1"
              local device_type="$2"  # "cache" or "vdev"

              if ! zpool status rpool | grep -q "$device_id"; then
                echo "ERROR: Device $device_id not found in pool rpool ($device_type)"
                return 1
              fi
              return 0
            }

            # Function to validate vdev configuration
            validate_vdevs() {
              local expected_mode="${config.storage.zfs.pool.mode}"
              local pool_status=$(zpool status rpool)

              # Check if pool exists
              if ! zpool list rpool >/dev/null 2>&1; then
                echo "ERROR: ZFS pool 'rpool' does not exist"
                return 1
              fi

              # Validate each configured vdev device
              ${lib.concatMapStringsSep "\n" (
                device: let
                  deviceStr =
                    if builtins.isString device
                    then device
                    else device.device;
                in ''
                  if ! check_device_in_pool "${deviceStr}" "vdev"; then
                    echo "ERROR: Vdev device ${deviceStr} not found in pool"
                    exit 1
                  fi
                ''
              )
              (lib.lists.flatten config.storage.zfs.pool.vdevs)}

              # Check pool mode matches configuration
              if ! echo "$pool_status" | grep -q "$expected_mode"; then
                echo "WARNING: Pool mode may not match expected configuration ($expected_mode)"
              fi

              echo "✓ All vdev devices validated successfully"
              return 0
            }

            # Function to validate cache configuration
            validate_cache() {
              ${lib.concatMapStringsSep "\n" (
              name: let
                device = config.storage.zfs.pool.cache.${name};
                deviceStr =
                  if builtins.isString device
                  then device
                  else device.device;
              in ''
                if ! check_device_in_pool "${deviceStr}" "cache"; then
                  echo "ERROR: Cache device ${deviceStr} (${name}) not found in pool"
                  exit 1
                fi
              ''
            ) (builtins.attrNames config.storage.zfs.pool.cache)}

              echo "✓ All cache devices validated successfully"
              return 0
            }

            # Run validations
            if validate_vdevs && validate_cache; then
              echo "✓ ZFS pool validation completed successfully"
            else
              echo "✗ ZFS pool validation failed"
              exit 1
            fi
          '';
          deps = ["zfs"];
        };

        # Script 2: Validate datasets and their options
        zfs-dataset-validation = {
          text = ''
            echo "Running ZFS dataset validation..."

            # Function to check if dataset exists
            check_dataset_exists() {
              local dataset="$1"
              if ! zfs list "$dataset" >/dev/null 2>&1; then
                echo "ERROR: Dataset $dataset does not exist"
                return 1
              fi
              return 0
            }

            # Function to validate dataset options
            validate_dataset_options() {
              local dataset="$1"
              local expected_options="$2"

              # Parse expected options (format: "option=value option2=value2")
              echo "$expected_options" | tr ' ' '\n' | while IFS='=' read -r option expected_value; do
                if [ -n "$option" ] && [ -n "$expected_value" ]; then
                  local actual_value=$(zfs get -H -o value "$option" "$dataset" 2>/dev/null)
                  if [ "$actual_value" != "$expected_value" ]; then
                    echo "ERROR: Dataset $dataset option $option is '$actual_value', expected '$expected_value'"
                    return 1
                  fi
                fi
              done
              return 0
            }

            # Validate root dataset
            echo "Validating root dataset..."
            if check_dataset_exists "rpool"; then
              root_options=""
              ${lib.concatMapStringsSep "\n" (
              option: let
                value = config.storage.zfs.rootDataset.${option};
              in
                lib.optionalString (value != null) ''
                  root_options="$root_options ${option}=${toString value}"
                ''
            ) ["canmount" "xattr" "acltype" "relatime" "compression" "encryption" "keyformat" "keylocation" "recordsize" "sync" "atime"]}

              # Add autoSnapshot option
              ${lib.optionalString (config.storage.zfs.rootDataset.autoSnapshot != null) ''
              root_options="$root_options com.sun:auto-snapshot=${
                if config.storage.zfs.rootDataset.autoSnapshot
                then "true"
                else "false"
              }"
            ''}

              if validate_dataset_options "rpool" "$root_options"; then
                echo "✓ Root dataset options validated"
              else
                echo "✗ Root dataset validation failed"
                exit 1
              fi
            else
              echo "✗ Root dataset validation failed"
              exit 1
            fi

            # Validate configured datasets
            ${lib.concatMapStringsSep "\n" (
              name: let
                dataset = config.storage.zfs.datasets.${name};
              in ''
                echo "Validating dataset: rpool/${name}"
                if check_dataset_exists "rpool/${name}"; then
                  dataset_options=""
                  ${lib.concatMapStringsSep "\n" (
                  option: let
                    value = dataset.${option};
                  in
                    lib.optionalString (value != null) ''
                      dataset_options="$dataset_options ${option}=${toString value}"
                    ''
                ) ["canmount" "xattr" "acltype" "relatime" "compression" "encryption" "keyformat" "keylocation" "recordsize" "sync" "atime"]}

                  # Add autoSnapshot option
                  ${lib.optionalString (dataset.autoSnapshot != null) ''
                  dataset_options="$dataset_options com.sun:auto-snapshot=${
                    if dataset.autoSnapshot
                    then "true"
                    else "false"
                  }"
                ''}

                  # Add custom options
                  ${lib.concatMapStringsSep "\n" (
                  optName: let
                    optValue = dataset.options.${optName};
                  in ''
                    dataset_options="$dataset_options ${optName}=${toString optValue}"
                  ''
                ) (builtins.attrNames (dataset.options or {}))}

                  if validate_dataset_options "rpool/${name}" "$dataset_options"; then
                    echo "✓ Dataset rpool/${name} options validated"
                  else
                    echo "✗ Dataset rpool/${name} validation failed"
                    exit 1
                  fi
                else
                  echo "✗ Dataset rpool/${name} validation failed"
                  exit 1
                fi
              ''
            ) (builtins.attrNames config.storage.zfs.datasets)}

            echo "✓ ZFS dataset validation completed successfully"
          '';
          deps = ["zfs" "zfs-pool-validation"];
        };
      };
    }
    (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;
          };
        };
      };
    })
  ]);
}