{ lib, ... }:

let
  rootFslabel = "external";
  rootFsDevice = "/dev/disk/by-label/${rootFslabel}";

  externalModule =
    partitionTableType:
    {
      config,
      lib,
      pkgs,
      ...
    }:
    {
      virtualisation.directBoot.enable = false;
      virtualisation.mountHostNixStore = false;
      virtualisation.useEFIBoot = partitionTableType == "efi";

      # This stops the qemu-vm module from overriding the fileSystems option
      # with virtualisation.fileSystems.
      virtualisation.fileSystems = lib.mkForce { };

      boot.loader.grub.enable = true;
      boot.loader.grub.efiSupport = partitionTableType == "efi";
      boot.loader.grub.efiInstallAsRemovable = partitionTableType == "efi";
      boot.loader.grub.device = if partitionTableType == "efi" then "nodev" else "/dev/vda";

      boot.growPartition = true;

      fileSystems = {
        "/".device = rootFsDevice;
        "/".fsType = "ext4";
      };

      # Needed for installing bootloader
      system.switch.enable = true;

      system.build.diskImage = import ../lib/make-disk-image.nix {
        inherit config lib pkgs;
        label = rootFslabel;
        inherit partitionTableType;
        format = "raw";
        bootSize = "128M";
        additionalSpace = "0M";
        copyChannel = false;
      };
    };
in
{
  name = "grow-partition";

  meta.maintainers = with lib.maintainers; [ arianvp ];

  nodes = {
    efi = externalModule "efi";
    legacy = externalModule "legacy";
    legacyGPT = externalModule "legacy+gpt";
    hybrid = externalModule "hybrid";
  };

  testScript =
    { nodes, ... }:
    lib.concatLines (
      lib.mapAttrsToList (name: node: ''
        import os
        import subprocess
        import tempfile
        import shutil

        tmp_disk_image = tempfile.NamedTemporaryFile()

        shutil.copyfile("${node.system.build.diskImage}/nixos.img", tmp_disk_image.name)

        subprocess.run([
          "${node.virtualisation.qemu.package}/bin/qemu-img",
          "resize",
          "-f",
          "raw",
          tmp_disk_image.name,
          "+32M",
        ])

        # Set NIX_DISK_IMAGE so that the qemu script finds the right disk image.
        os.environ['NIX_DISK_IMAGE'] = tmp_disk_image.name

        ${name}.wait_for_unit("growpart.service")
        systemd_growpart_logs = ${name}.succeed("journalctl --boot --unit growpart.service")
        assert "CHANGED" in systemd_growpart_logs
        ${name}.succeed("systemctl restart growpart.service")
        systemd_growpart_logs = ${name}.succeed("journalctl --boot --unit growpart.service")
        assert "NOCHANGE" in systemd_growpart_logs

      '') nodes
    );
}
