PortValhalla/lib/modules/partition/disks.nix

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

  fs = import ./fs.nix;

  diskListVarName = "myDisks";
  isSwap = partition: builtins.elem partition.type [ fs.swap 19 ];

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

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

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

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

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

      config = let
        diskVarName = "${diskListVarName}[${config.id}]";
        diskVar = "\${${diskVarName}}";

        diskSelector = ''
          result="$(mktemp)"
          fish ${./choose-disk.fish} "$result" "Which disk do you wish to install the OS on?" ${./select.fish}
          ${diskVarName}="$(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 ""
            } ${diskVar}"
          }";
        wipeScript = script: fdiskScript script [] false;
        appendScript = index: script: fdiskScript script ["-N" (builtins.toString index)] true;

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

        fdiskCommands = lib.strings.concatLines
          (lib.optionals config.wipe [
            cleanup
            (wipeScript "label: gpt")
          ] ++ (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.swap} = "mkswap ${partVar}";
                  ${fs.ntfs} = "mkfs.ntfs -F ${partVar}";
                  ${fs.fat32} = "mkfs.fat -F 32 ${partVar}";
                };

                labelScripts = {
                  ${fs.ext4} = label: "e2label ${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
                  "sudo ${formatScripts.${format}}"
                ];

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

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

        deviceScript = if osDisk && config.devicePath == null then ''
          ${diskSelector}
        '' else ''
          ${diskVarName}=${config.devicePath}
          ${if osDisk then ''
            if [ ! -b ${diskVar} ]; then
              function fallback() {
                echo "Couldn't find the specified disk \"${diskVar}\"."
                if fish ${./confirm.fish} "Do you want to install the OS on another disk?"; then
                  ${diskSelector}
                else
                  exit 1
                fi
              }

              fallback
            fi
          '' else
            ""}
        '';

        script = 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.enum (builtins.attrValues fs);
          description = "The file system format of the partition.";
          default = if (isSwap config) then
            fs.swap
          else
            throw ("Partition format not specified.");
        };

        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 = {
      partition = {
        rootDir = mkOption {
          type = types.str;
          description = "The root of the installation directory to mount disks into.";
          default = "/mnt";
        };

        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.str;
          description = "The script for partitioning the system's disks.";
        };
      };
    };
  };

  config = {
    valhalla = {
      partition = {
        script = lib.mkDefault (let
          cfg = config.valhalla.partition;
          inherit (cfg) os rootDir;
          inherit (lib.strings) normalizePath;
          partPath = part: "/dev/disk/by-label/${part.label}";
          disks = ([ os ] ++ (builtins.attrValues cfg.disks));
          partitions = (builtins.concatMap (_: (builtins.attrValues _.partitions)) disks);

          mountScript = lib.strings.concatLines (builtins.concatMap (
            _: [
              probeScript
              (builtins.concatStringsSep " " ([
                  "sudo"
                  "mount"
                  "--mkdir"
                ] ++ (lib.optionals (_.format == "ntfs") [
                  "-t" "ntfs3"
                ]) ++ [
                  (builtins.concatStringsSep " " (builtins.map (_: "-o ${_}") _.mountOptions))
                  (partPath _)
                  (normalizePath "/${rootDir}/${_.mountPoint}")
                ]))
            ]) (lib.lists.sortOn
              (_: normalizePath "/${_.mountPoint}")
                (builtins.filter (_: _.mountPoint != null) partitions)));

          swapScript = lib.strings.concatLines (builtins.map (
            _: ''
              ${probeScript}
              sudo swapon ${partPath _}
            '') (builtins.filter (_: _.useSwap) partitions));
        in lib.strings.concatLines ([
          "#!/bin/bash"
          "set -o errexit"
        ]
          ++ (builtins.map (_: _.deviceScript) disks)
          ++ lib.optionals ((builtins.length disks) > 0) [
            ''echo "$(tput setaf 3)==== WARNING ====$(tput sgr0)"''
            (''echo "Continuing this script will alter the partitions of ''
              + (lib.strings.concatStringsSep ", " (builtins.map (_: "${_.deviceVariable}") (lib.lists.init disks)))
              + (if (builtins.length disks) > 1 then " and " else "") + (lib.lists.last disks).deviceVariable + ''"'')
            ''
              if ! fish ${./confirm.fish} "Are you sure you want to continue?" "n"; then
                exit 1
              fi
            ''
          ] ++ (builtins.map (_: _.script) disks) ++ [
            mountScript
            swapScript
          ]));
      };
    };
  };
}