PortValhalla/lib/modules/fileSystems/disks.nix

{ lib, config, ... }:
let
  inherit (lib) types mkOption;

  fs = import ./fs.nix;

  cfg = config.valhalla.fileSystems;
  deviceListVarName = "myDevices";
  isSwap = partition: builtins.elem partition.type [ fs.swap 19 ];

  probeScript = builtins.concatStringsSep "\n" [
    "partprobe 2> /dev/null || true"
    "udevadm trigger"
  ];

  mkDeviceType = types.submodule (
    { config, name, ... }: {
      options = {
        id = mkOption {
          type = types.str;
          description = "The internal identifier of the device.";
          internal = true;
        };

        wipe = mkOption {
          type = types.bool;
          description = "A value indicating whether the device should be wiped.";
          default = !(lib.lists.any (_: _.keepExisting) (builtins.attrValues config.partitions));
        };

        name = mkOption {
          type = types.nullOr types.str;
          description = "The name of the device.";
          default = name;
        };

        path = mkOption {
          type = types.nullOr types.str;
          description = "The path to the device.";
          default =
            if config.name == null then
              null
            else
              "/dev/${config.name}";
        };

        deviceScript = mkOption {
          type = types.str;
          description = "A command for loading the device path into the device variable";
          internal = true;
        };

        deviceVariable = mkOption {
          type = types.str;
          description = "The name of the variable holding the name of the disk";
          internal = true;
        };

        partitions = mkOption {
          type = types.attrsOf (types.nullOr partitionType);
          description = "The partitions of the disk.";
          default = { };
        };

        scripts = {
          init = mkOption {
            type = types.str;
            description = "A script for loading the device path into the device variable";
          };

          partition = mkOption {
            type = types.str;
            description = "A script for partitioning and formatting the device.";
          };
        };
      };

      config =
        let
          deviceVarName = "${deviceListVarName}[${config.id}]";
          deviceVar = "\${${deviceVarName}}";

          deviceSelector = ''
            result="$(mktemp)"
            fish ${./choose-device.fish} "$result" "Please select the \"${name}\" device:" ${./select.fish}
            ${deviceVarName}="$(cat "$result")"
          '';

          partitions = lib.lists.sortOn (_: _.index)
            (builtins.filter (_: _ != null)
              (builtins.attrValues config.partitions));

          mkType = type:
            lib.strings.escapeShellArg (
              if builtins.isInt type then
                "${lib.trivial.toHexString type}"
              else
                type
            );

          fdiskCommand = arguments: "sudo sfdisk ${arguments}";
          fdiskScript = script: args: append:
            "echo ${script} | ${
            fdiskCommand "${builtins.concatStringsSep " " args} ${
              if append then "--append" else ""
            } ${deviceVar}"
          }";
          appendScript = index: script: fdiskScript script [ "-N" (builtins.toString index) ] true;

          cleanup = lib.strings.concatLines (builtins.map
            (partition: "${fdiskCommand "--delete ${deviceVar} ${toString partition.index}"} || true")
            (lib.lists.sortOn
              (partition: partition.index * -1)
              (builtins.filter (_: !_.keepExisting) partitions)));

          fdiskCommands = lib.strings.concatLines
            (lib.optionals config.wipe [
              cleanup
              (fdiskScript "label: gpt" [ ] false)
            ] ++ (builtins.concatMap
              (
                partition:
                let
                  inherit (partition) format index keepExisting label sizeScript type;

                  partVarName = "myPartition";
                  partVar = "\${${partVarName}}";

                  sizeOption = ''
                    ${sizeScript} | sed -e "s/.*[^[:space:]]/size=\0/"
                  '';

                  formatScripts = {
                    ${fs.ext4} = "mkfs.ext4 -F ${partVar}";
                    ${fs.btrfs} = "mkfs.btrfs --force ${partVar}";
                    ${fs.swap} = "mkswap ${partVar}";
                    ${fs.ntfs} = "mkfs.ntfs -F ${partVar}";
                    ${fs.fat32} = "mkfs.fat -F 32 ${partVar}";
                  };

                  labelScripts = {
                    ${fs.ext4} = label: "e2label ${partVar} ${label}";
                    ${fs.btrfs} = label: "btrfs filesystem label ${partVar} ${label}";
                    ${fs.swap} = label: "swaplabel ${partVar} --label ${label}";
                    ${fs.ntfs} = label: "ntfslabel ${partVar} ${label}";
                    ${fs.fat32} = label: "fatlabel ${partVar} ${label}";
                  };

                  create = lib.strings.concatLines ([
                    (appendScript index ''${toString index}: "$(${sizeOption})" type=${mkType type}'')
                    probeScript
                  ] ++ (lib.optionals (format != null) [
                    "sudo ${formatScripts.${format}}"
                  ]));

                  fallback = ''
                    if ! { ls "${partVar}" 2>&1; } > /dev/null
                    then
                      ${create}
                    fi
                  '';
                in
                [
                  ''local diskPath="$(find -L /dev/disk/by-diskseq -samefile ${deviceVar})"''
                  ''local ${partVarName}="$diskPath-part${toString index}"''
                  (if keepExisting then fallback else create)
                ] ++ (lib.optionals (format != null) [
                  "sudo ${labelScripts.${format} label}"
                ])
              )
              partitions));

          fixType = lib.strings.concatLines (builtins.concatMap
            (
              partition:
              lib.optional
                (partition.keepExisting && !(builtins.isNull partition.type))
                ''sudo sfdisk --part-type ${deviceVar} ${toString partition.index} ${mkType partition.type}''
            )
            partitions);
        in
        {
          id = "disk-${name}";
          deviceVariable = deviceVar;

          scripts = {
            init =
              if config.path == null then ''
                ${deviceSelector}
              '' else ''
                ${deviceVarName}=${config.path}
                if [ ! -b ${deviceVar} ]; then
                  function fallback() {
                    echo "Couldn't find the specified disk \"${deviceVar}\"."
                    if fish ${./confirm.fish} "Do you want to choose a different \"${name}\" disk?"; then
                      ${deviceSelector}
                    else
                      exit 1
                    fi
                  }

                  fallback
                fi
              '';

            partition = lib.mkDefault ''
              function partition() {
                ${if (!config.wipe) then cleanup else ""}
                ${probeScript}
                ${fdiskCommands}
                ${fixType}
              }

              partition
            '';
          };
        };
    }
  );

  partitionType = types.submodule (
    { name, config, ... }: {
      options = {
        index = mkOption {
          type = types.int;
          description = "The index of the partition.";
        };

        label = mkOption {
          type = types.str;
          description = "The label of the partition.";
          default = name;
        };

        keepExisting = mkOption {
          type = types.bool;
          description = "A value indicating whether the partition should be left untouched if it already exists.";
          default = false;
        };

        type = mkOption {
          type = types.nullOr (types.either types.str types.int);
          description = "The type of the partition.";
          default = null;
        };

        format = mkOption {
          type = types.nullOr (types.enum (builtins.attrValues fs));
          description = "The file system format of the partition.";
          default =
            if (isSwap config) then
              fs.swap
            else
              null;
        };

        size = mkOption {
          type = types.nullOr types.str;
          description = "The size of the partition.";
          default = null;
        };

        sizeScript = mkOption {
          type = types.str;
          description = "A script for printing the size to the console.";
          internal = true;
        };

        useSwap = mkOption {
          type = types.bool;
          description = "A value indicating whether this partition should be used as swap.";
          default = isSwap config;
        };

        mountPoint = mkOption {
          type = types.nullOr types.str;
          description = "The mountpoint of the partition.";
          default = null;
        };

        mountOptions = mkOption {
          type = types.listOf types.str;
          description = "The options to apply to the mount.";
          default = [ ];
        };
      };

      config = {
        sizeScript = (if isSwap config then
          ''echo "$(cat /proc/meminfo | awk -F " " '/^MemTotal/ { print $2 }' | awk '{ print int((($1 / 1024 / 1024) * 0.75) + 0.5)}')"G''
        else
          "echo ${lib.strings.escapeShellArg (toString config.size)}");
      };
    }
  );
in
{
  options = {
    valhalla = {
      fileSystems = {
        diskSetup = {
          devices = mkOption {
            type = types.attrsOf (mkDeviceType);
            description = "The disk devices to format.";
            default = { };
          };

          scripts = {
            init = mkOption {
              type = types.str;
              description = "The script for initializing the disk partitioning script.";
            };

            partition = mkOption {
              type = types.str;
              description = "The script for partitioning the disks.";
            };

            swap = mkOption {
              type = types.str;
              description = "The script for enabling swap devices.";
            };
          };
        };
      };
    };
  };

  config = {
    valhalla = {
      fileSystems = {
        mounts = (lib.attrsets.concatMapAttrs
          (
            name: device:
              lib.attrsets.concatMapAttrs
                (
                  name: partition:
                    if partition.mountPoint != null then {
                      ${partition.mountPoint} = {
                        device = "/dev/disk/by-label/${partition.label}";
                        fsType = partition.format;
                        options = partition.mountOptions;
                      };
                    } else { }
                )
                device.partitions
          )
          cfg.diskSetup.devices);

        diskSetup = {
          scripts =
            let
              partPath = part: "/dev/disk/by-label/${part.label}";
              disks = ((builtins.attrValues cfg.diskSetup.devices));
              partitions = (builtins.concatMap (_: (builtins.attrValues _.partitions)) disks);
            in
            {
              init = lib.strings.concatLines (builtins.map (_: _.scripts.init) disks);
              partition = lib.strings.concatLines (builtins.map (_: _.scripts.partition) disks);
              swap = lib.strings.concatLines (
                (builtins.map
                  (
                    _: ''
                      ${probeScript}
                      sudo swapon ${partPath _}
                    ''
                  )
                  (builtins.filter (_: _.useSwap) partitions))
              );
            };
        };
      };
    };
  };
}