{ config, pkgs, lib, ... }:
with lib;
let
  cfg = config.containers;

  concatMapLines = concatMapStringsSep "\n";

  wait_cmds = concatMapLines
    (dataset: ''until zfs list "${dataset}"; do echo "Waiting for dataset \"${dataset}\"..."; sleep 0.5; done'');

  zoned_cmds = concatMapLines (dataset: ''zfs set zoned=on "${dataset}"'');
  delegate_cmds = concatMapLines (dataset: ''zfs zone /proc/"$LEADER_PID"/ns/user "${dataset}"'');
in
{
  options.containers = mkOption {
    type = types.attrsOf (types.submodule (
      { config, ... }:
      {
        options.zfs_datasets = mkOption {
          description = "ZFS datasets to mount into the container";
          type = types.listOf types.str;
          default = [];
        };

        config = lib.mkIf (length config.zfs_datasets >= 1) {
          allowedDevices = [
            {
              node = "/dev/zfs";
              modifier = "rw";
            }
          ];

          bindMounts."/dev/zfs" = {
            hostPath = "/dev/zfs";
            isReadOnly = false;
          };

          # This makes systemd-nspawn immediately notify the parent service that the container is
          # ready without waiting for its services to start. ExecStartPost is run when this ready
          # signal is received. Turning this off creates a deadlock where the container's
          # `zfs-datasets` service needs the `ExecStartPost` script to run, but that script won't
          # run until the service exits
          extraFlags = [ "--notify-ready=no" ];

          config = { pkgs, ... }: {
            systemd.targets."zfs" = {
              wantedBy = [ "multi-user.target" ];
            };

            systemd.services."zfs-datasets" = {
              description = "Waits for container ZFS datasets to become available";

              path = with pkgs; [
                zfs
              ];

              script = ''
                  ${wait_cmds config.zfs_datasets}
                  echo "All ZFS datasets available"
              '';

              serviceConfig = {
                Type = "oneshot";
                RemainAfterExit = true;
                TimeoutStartSec = 60 * 5;
              };

              wantedBy = [ "zfs.target" ];
            };
          };
        };
      }
    ));
  };

  config.systemd.services = mapAttrs' (name: config:
    {
      name = "container@${name}";

      value = lib.mkIf (length config.zfs_datasets >= 1) {
        path = with pkgs; [
          zfs
          systemd
        ];

        preStart = mkBefore ''
          ${zoned_cmds config.zfs_datasets}
          echo "Enabled zoned setting on zfs datasets"
        '';

        postStart = mkBefore ''
          LEADER_PID=$(machinectl show -P Leader ${name})
          ${delegate_cmds config.zfs_datasets}
          echo "Delegated zfs dataset to namespace of PID $LEADER_PID"
        '';
      };
  }) cfg;
}