{ config, lib, ... }:

with lib;

# default pool for root device is named rpool, cannot be changed atm.

# IDEA
# mine.disks.root = "/dev/sda1";
# mine.disks.pools.rpool.device = "/dev/disk/....";
# mine.disks.pools.rpool.createDisk = false;
# mine.disks.pools.rpool.extra = {};
# mine.disks.pools.rpool.datasets = { "local/nix" = { mountpoint = "/nix/store"; extra = {}; }; };

let
  cfg = config.mine.disks;

  datasetOpts = { name, config, ... }: {
    options = {
      name = lib.mkOption {
        type = types.str;
        default = name;
        description = "The name of the name dataset";
      };

      mountpoint = lib.mkOption {
        type = types.nullOr types.str;
        description = "Where to mount the ZFS dataset";
      };

      extra = lib.mkOption {
        type = types.attrsOf types.anything;
        default = {};
      };
    };
  };

  poolOpts = { name, config, ... }: {
    options = {
      name = lib.mkOption {
        type = types.str;
        default = name;
        description = "The name of the zfs pool, defaults to name";
      };

      diskoDiskName = lib.mkOption {
        type = types.str;
        default = "disk-${name}";
        description = "The name of the disk attribute in disko, defaults to disk-<name>";
      };

      device = lib.mkOption {
        type = types.nullOr types.str;
        default = null;
        description = "Device to use for the pool (needed if the pool should be created)";
      };

      createDisk = lib.mkOption {
        type = types.bool;
        default = true;
        description = "Indicates if the disk should be created (only set to false for the root)";
      };

      extra = lib.mkOption {
        type = types.attrsOf types.anything;
        default = {};
      };

      datasets = mkOption {
        type = types.attrsOf (types.submodule datasetOpts);
        default = {};
      };
    };
  };

  sources = import ./../sources/sources.nix;
in {
  imports = [
    (sources.disko + "/module.nix")
  ];

  options.mine.disks = {
    disk = mkOption {
      type = types.str;
      default = "/dev/sda";
    };

    rollbackOnBoot = mkOption {
      type = types.bool;
      default = true;
    };

    pools = mkOption {
      type = types.attrsOf (types.submodule poolOpts);
      default = {};
    };
  };

  config = {
    # default rpool (cannot be changed)
    mine.disks.pools.rpool = {
      createDisk = false;

      datasets = {
        "local/nix".mountpoint = "/nix";
        "root" = { mountpoint = "/"; extra = { postCreateHook = "zfs snapshot rpool/root@blank"; }; };
        "safe/persistent".mountpoint = "/state/root";
      };
    };

    # state things are needed on BOOT!
    fileSystems."${cfg.pools.rpool.datasets."safe/persistent".mountpoint}".neededForBoot = true;

    # rollback on boot!
    boot.initrd.postDeviceCommands = lib.mkIf cfg.rollbackOnBoot (lib.mkAfter ''
      zfs rollback -r rpool/root@blank
    '');
    
    disko.devices = {
      disk = let
        diskPoolsToCreate = lib.filterAttrs (_: v: v.createDisk) cfg.pools;
        extraDisks = lib.mapAttrs' (_: v: lib.nameValuePair v.diskoDiskName (lib.recursiveUpdate {
          type = "disk";
          device = v.device;
          content = {
            type = "gpt";
            partitions = {
              zfs = {
                size = "100%";
                content = {
                  type = "zfs";
                  pool = v.name;
                };
              };
            };
          };
        } v.extra)) diskPoolsToCreate;
      in {
        root = {
          type = "disk";
          device = cfg.disk;
          content = {
            type = "gpt";
            partitions = {
              boot = {
                name = "boot";
                size = "1M";
                type = "EF02";
              };
              esp = {
                name = "ESP";
                size = "500M";
                type = "EF00";
                content = {
                  type = "filesystem";
                  format = "vfat";
                  mountpoint = "/boot";
                };
              };
  
              luks = {
                size = "100%";
                content = {
                  type = "luks";
                  name = "cryptroot";
                  settings.allowDiscards = true;
  
                  content = {
                    type = "zfs";
                    pool = "rpool";
                  };
                };
              };
            };
          };
        };
      } // extraDisks;
  
      zpool = let
        # helper functions
        makeZFSDatasets = datasets: (lib.mapAttrs' (n: v: lib.nameValuePair v.name (lib.recursiveUpdate {
            type = "zfs_fs";
            mountpoint = v.mountpoint;
            options.mountpoint = "legacy";
        } (if v ? extra then v.extra else {}))) datasets);

        mkZFSPool = zfsPool: {
          type = "zpool";
          # rootFsOptions.compression = "zstd";
          rootFsOptions = {
            compression = "on";
            atime = "off";
            acltype = "posixacl";
            xattr = "sa";
          };

          datasets = makeZFSDatasets zfsPool.datasets;
        };
      in lib.mapAttrs' (n: v: lib.nameValuePair v.name (mkZFSPool v)) cfg.pools;
      #   {
      #   rpool = {
      #     type = "zpool";
      #     # rootFsOptions.compression = "zstd";
      #     rootFsOptions = {
      #       compression = "on";
      #       atime = "off";
      #       acltype = "posixacl";
      #       xattr = "sa";
  
      #       # test
      #       # "com.sun:auto-snapshot" = "false";
      #       # "com.klarasystems:vdev_zaps_v2" = "false";
      #     };
  
      #     datasets = let
      #       baseDatasets = {
      #         "/" = { dataset = "root"; extra = { postCreateHook = "zfs snapshot rpool/root@blank"; }; };
      #         "/nix".dataset = "local/nix";
      #         "/state/stash".dataset = "local/stash";
      #         "/state/home".dataset = "safe/home";
      #         "/state/root".dataset = "safe/persistent";
  
      #         # extra datasets
      #         "/srv/forgejo" = { dataset = "safe/svcs/forgejo"; extra.options.quota = "5G"; };
      #       };
      #     in (makeZFSDatasets baseDatasets);
      #   };
      # };
    };
  
    # fileSystems."/state/root".neededForBoot = true;
    # fileSystems."/state/home".neededForBoot = true;
    # fileSystems."/state/stash".neededForBoot = true;
  };
}