nix-config/modules/nixos/disko/default.nix

{
  config,
  lib,
  namespace,
  ...
}:
with lib;
let
  cfg = config.${namespace}.hardware.disko;

  defaultBtrfsMountOptions = [
    "compress=${cfg.compression}"
    "noatime"
  ];
  defaultBcachefsMountOptions = [
    "noatime"
  ];

  subvolumes =
    let
      make =
        name: subvolume:
        nameValuePair "${name}" {
          mountOptions =
            if subvolume.mountOptions == null then
              if cfg.filesystem == "btrfs" then defaultBtrfsMountOptions else defaultBcachefsMountOptions
            else
              subvolume.mountOptions;
          mountpoint = if subvolume.mountPoint == null then "/${name}" else subvolume.mountPoint;
        };
    in
    mapAttrs' make cfg.subvolumes;

  # BTRFS root partition configuration
  root = {
    name = "${config.${namespace}.network.hostName}-${cfg.filesystem}-root";
    size = "100%";
    content = {
      type = cfg.filesystem;
    }
    // (
      if cfg.filesystem == "btrfs" then
        {
          extraArgs = [ "-f" ]; # Override existing partition
          # Subvolumes must set a mountpoint in order to be mounted,
          # unless their parent is mounted
          subvolumes = subvolumes;
        }
      else
        {
          # This refers to a filesystem in the `bcachefs_filesystems` attrset below.
          filesystem = "mounted_subvolumes_in_multi";
          label = "ssd.ssd1";
          extraFormatArgs = [
            "--discard"
          ];
        }
    );
  };

  # Luks root partition configuration
  luksRoot = {
    name = "${config.${namespace}.network.hostName}-cryptroot";
    size = "100%";
    content = {
      type = "luks";
      name = "${config.${namespace}.network.hostName}-cryptroot";
      extraOpenArgs = [
        "--allow-discards"
        "--perf-no_read_workqueue"
        "--perf-no_write_workqueue"
      ];
      settings = {
        crypttabExtraOpts = [
          "tpm2-device=auto"
          "fido2-device=auto"
          "token-timeout=10"
        ];
      };
      content = {
        type = cfg.filesystem;
      }
      // (
        if cfg.filesystem == "btrfs" then
          {
            extraArgs = [ "-f" ]; # Override existing partition
            # Subvolumes must set a mountpoint in order to be mounted,
            # unless their parent is mounted
            subvolumes = subvolumes;
          }
        else
          {
            # This refers to a filesystem in the `bcachefs_filesystems` attrset below.
            filesystem = "mounted_subvolumes_in_multi";
            label = "ssd.ssd1";
            extraFormatArgs = [
              "--discard"
            ];
          }
      );
    };
  };
in
{
  imports = [ ./options.nix ];
  config = lib.mkIf cfg.enable {
    assertions = [
      {
        assertion = cfg.rootDisk != "";
        message = "mjallen.hardware.disko.rootDisk must be set to a non-empty device path (e.g. \"/dev/nvme0n1\").";
      }
      {
        assertion = !(cfg.enableSwap && cfg.swapSize == "");
        message = "mjallen.hardware.disko.swapSize must be a non-empty size string when enableSwap is true (e.g. \"16G\").";
      }
      {
        assertion = cfg.compression != "";
        message = "mjallen.hardware.disko.compression must be a non-empty compression type (e.g. \"zstd\").";
      }
    ];

    disko.devices = lib.mkMerge [
      {
        nodev."/" = {
          fsType = "tmpfs";
          mountOptions = [
            "mode=755"
            "defaults"
            "size=25%"
          ];
        };
        disk = {
          main = {
            device = cfg.rootDisk;
            type = "disk";
            imageSize = "32G";
            content = {
              type = "gpt";
              partitions = {
                firmware = lib.mkIf cfg.firmware.enableFirmware {
                  priority = 1;
                  type = "0700";
                  name = "${config.${namespace}.network.hostName}-FIRMWARE";
                  start = "1M";
                  end = "1G";
                  content = {
                    type = "filesystem";
                    format = "vfat";
                    mountpoint = "/boot/firmware";
                    mountOptions = [ "umask=0077" ];
                  };
                };
                ESP = {
                  priority = if cfg.firmware.enableFirmware then 2 else 1;
                  type = "EF00";
                  size = "500M";
                  name = "${config.${namespace}.network.hostName}-EFI";
                  content = {
                    type = "filesystem";
                    format = "vfat";
                    mountpoint = "/boot";
                    mountOptions = [ "umask=0077" ];
                  };
                };

                swap = lib.mkIf cfg.enableSwap {
                  type = "8200";
                  size = cfg.swapSize;
                };

                root = if cfg.enableLuks then luksRoot else root;
              };
            };
          };
          # firmware = lib.mkIf cfg.firmware.enableFirmware {
          #   device = cfg.firmware.firmwareDisk;
          #   type = "disk";
          #   imageSize = "1G";
          #   content = {
          #     type = "table";
          #     format = "msdos";
          #     partitions = [
          #       {
          #         name = "${config.${namespace}.network.hostName}-FIRMWARE";
          #         start = "1M";
          #         end = "1G";
          #         content = {
          #           type = "filesystem";
          #           format = "vfat";
          #           mountpoint = "/boot/firmware";
          #           mountOptions = [ "umask=0077" ];
          #         };
          #       }
          #     ];
          #   };
          # };
        };

        # configure Bcachefs
        bcachefs_filesystems = lib.mkIf (cfg.filesystem == "bcachefs") {
          mounted_subvolumes_in_multi = {
            type = "bcachefs_filesystem";
            # passwordFile = "/etc/nixos/test.key";
            extraFormatArgs = [
              "--compression=${cfg.compression}"
            ];
            subvolumes = subvolumes;
          };
        };
      }
    ];
  };
}