PortValhalla/lib/modules/partition.nix

{ lib, config, ... }:
  let
    inherit (lib) types mkOption;
    diskVarName = "myDisk";
    diskVar = ''''${${diskVarName}}'';
    isSwap = partition: builtins.elem partition.type ["swap" 19];

    diskSelector = ''
      . ${./../../scripts/Common/Scripts/choose-disk.sh};
      chooseDisk ${diskVarName} "Which disk do you wish to install the OS on?";
    '';

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

          deviceName = mkOption {
            type = types.nullOr types.str;
            description = "The name of the device.";
            default = if osDisk then null else name;
          };

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

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

          script = mkOption {
            type = types.str;
            description = "The script for formatting the disk.";
            internal = true;
          };
        };

        config =
        let
          partitions = builtins.foldl'
            (list: predicate: lib.lists.sortOn predicate list)
            (builtins.attrValues config.partitions)
            [
              (_: _.index)
              (_: _.order)
              (_: if ((!isSwap _) && (_.size == null)) then 1 else 0)
            ];

          fdiskCommand = arguments: "sudo sfdisk ${arguments}";
          fdiskScript = script: append: "echo ${script} | ${fdiskCommand "${if append then "--append" else ""} ${diskVar}"}";
          wipeScript = script: fdiskScript script false;
          appendScript = script: fdiskScript script true;

          cleanup = lib.strings.concatLines (builtins.map
            (partition: fdiskCommand "--delete ${diskVar} ${toString partition.index}")
            (builtins.filter (_: !_.keepExisting) partitions));

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

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

                  create = appendScript "${toString index}: \"$(${sizeOption})\" type=${lib.strings.escapeShellArg type}";
                  fallback = ''
                    if ! { ls "${diskVar}"?(p)${toString partition.index} 2>&1; } > /dev/null
                    then
                      ${create}
                    fi
                  '';
                in [
                  (if keepExisting then fallback else create)
                ])
              partitions)
          );

          fixType = lib.strings.concatLines (
            builtins.concatMap
              (partition: lib.optional
                partition.keepExisting
                ''sudo sfdisk --part-type ${diskVar} ${toString partition.index} ${lib.strings.escapeShellArg partition.type}'')
              partitions);
        in {
          script = ''
            function partition() {
              ${
                if osDisk && config.devicePath == null then ''
                  ${diskSelector}
                ''
                else
                  "local ${diskVarName}=${config.devicePath}"}

              ${if (!config.wipe) then cleanup else ""}

              ${fdiskCommands}
              ${fixType}
            }

            partition
          '';
        };
      });

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

          order = mkOption {
            type = types.int;
            description = "The sort order of the partition creation.";
            default = if (!(isSwap config)) && ((builtins.match ".*[^[:space:]].*" (toString config.size)) == null) then 1 else 0;
          };

          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;
          };

          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;
          };

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

        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 = {
      partition = {
        os = mkOption {
          type = mkDiskType true;
          description = "The partition layout of the OS disk.";
        };

        disks = mkOption {
          type = types.attrsOf (mkDiskType false);
          description = "The additional disks to format.";
          default = {};
        };

        script = mkOption {
          type = types.lines;
          description = "The script for partitioning the system's disks.";
          internal = true;
        };
      };
    };

    config = {
      partition = {
        os = {
          partitions = {
            Boot = {
              index = 1;
              type = "uefi";
              size = "+1G";
            };

            Swap = {
              index = 2;
              type = "swap";
            };

            OS = {
              index = 3;
              type = "linux";
              size = "";
            };
          };
        };

        script =
          let
            inherit (config.partition)
              os
              disks
            ;
          in lib.strings.concatLines (builtins.map (_: _.script) ([os] ++ (builtins.attrValues disks)));
      };
    };
  }