{ 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
          ])
        );
      };
    };
  };
}