{ config, lib, ... }:
let
  cfg = config.jalr.luksUsbUnlock;
in
{
  options.jalr.luksUsbUnlock = with lib; with lib.types; {
    enable = mkEnableOption "unlock LUKS volumes with a USB device on boot";
    devices = mkOption {
      default = { };
      example = {
        cryptroot = {
          keyPath = "/path/to/the/key";
          usbDevice = "by-label/MY_USB";
        };
      };
      type = types.attrsOf (types.submodule {
        options = {
          keyPath = mkOption {
            example = "/mykey.key";
            description = mdDoc ''
              Path to the key file inside the USB device's filesystem.
              `/` is relative to the device's filesystem root.
            '';
            type = types.str;
          };

          usbDevice = mkOption {
            example = "by-label/BOOTKEY";
            description = mdDoc ''
              Path to the USB device that contains the keys. (Path relative to `/dev/disk/`)
            '';
            type = types.str;
          };

          waitForDevice = mkOption {
            default = 5;
            example = 10;
            description = mdDoc ''
              How many seconds to wait for the USB device to be detected by the
              kernel.
            '';
            type = types.ints.unsigned;
          };
        };
      });
    };
  };
  config = lib.mkIf cfg.enable (
    let
      makeUsbDevPath = usbDevice: "/dev/disk/" + usbDevice;
      makeMountPath = usbDevice: "/key/" + (builtins.hashString "md5" usbDevice);
      usbFsType = "vfat";

      mapAttrsNameValue = f: set:
        lib.listToAttrs (map f (lib.attrsToList set));
    in
    {
      boot.initrd = {
        kernelModules = [ "uas" "usbcore" "usb_storage" "vfat" "nls_cp437" "nls_iso8859_1" ];
        systemd.services =
          let
            makeService = name: { usbDevice, waitForDevice, ... }:
              let
                usbDevPath = makeUsbDevPath usbDevice;
                usbMountPath = makeMountPath usbDevice;
              in
              {
                description = "Mount ${name} key";
                wantedBy = [ "cryptsetup.target" ];
                before = [ "systemd-cryptsetup@${name}.service" ];
                after = [ "systemd-modules-load.service" ];
                unitConfig.DefaultDependencies = "no";
                serviceConfig.Type = "oneshot";

                script = ''
                  if awk -v mountpoint="${usbMountPath}" '$2==mountpoint {f=1} END {exit !f}' /proc/mounts; then
                    exit 0
                  fi

                  attempts=0
                  while [ ! -e ${lib.escapeShellArg usbDevPath} ]; do
                    sleep 1
                    if [ $attempts -ge ${toString waitForDevice} ]; then
                      break;
                    fi
                    attempts=$((attempts+1))
                  done

                  if [ -e ${lib.escapeShellArg usbDevPath} ]; then
                    mkdir -m0500 -p ${lib.escapeShellArg usbMountPath}
                    mount \
                      -n \
                      -t ${lib.escapeShellArg usbFsType} \
                      -o ro,fmask=0137,dmask=0027 \
                      ${lib.escapeShellArg usbDevPath} \
                      ${lib.escapeShellArg usbMountPath}
                  fi
                '';
              };
          in
          mapAttrsNameValue
            ({ name, value }: {
              name = "luksusb-${name}";
              value = makeService name value;
            })
            cfg.devices;

        luks.devices = builtins.mapAttrs
          (_: { keyPath, usbDevice, ... }:
            let
              usbMountPath = makeMountPath usbDevice;
            in
            {
              keyFile = "${usbMountPath}/${keyPath}";
              keyFileTimeout = 1;
            })
          cfg.devices;
      };
    }
  );
}