# NOTE: ... is needed because dikso passes diskoFile
{
  lib,
  disk ? [ "/dev/vda" ],
  withSwap ? false,
  swapSize,
  configVars,
  ...
}:
### inspiration: https://github.com/chewblacka/nixos/blob/main/disko-config.nix
let 
  number_of_disks = if (builtins.length disks < 3) 
                    then builtins.length disks 
                    else throw "Error. Too many disks passed to disko.";
in
{
  disko.devices = {
    disk = {
      vda = {
        type = "disk";
        device = builtins.elemAt disks 0;
        content = {
          type = "gpt";
          partitions = {
            ESP = {
              # name = "ESP";
              priority = 1;
              start = "1MiB";
              end = "1G";
              type = "EF00";
              content = {
                type = "filesystem";
                format = "vfat";
                mountpoint = "/boot";
                mountOptions = [ "defaults" ]
              };
            };
            root = {
              # name = "root";
              start = "9G";
              end = "100%";
              content = {
                type = "btrfs";
                extraArgs = [ "-f" ]; # Override existing partition
                subvolumes = 
                  if (number_of_disks == 1) then 
                    { 
                      "@" = { };
                      "@/root" = {
                        mountpoint = "/";
                        mountOptions = [ "compress=zstd" "noatime" ];
                      };
                      "@/home" = {
                        mountpoint = "/home";
                        mountOptions = [ "compress=zstd" ];
                      };
                      "@/nix" = {
                        mountpoint = "/nix";
                        mountOptions = [ "compress=zstd" "noatime" ];
                      };
                      "@persist" = {
                        mountpoint = "${configVars.persistFolder}";
                        mountOptions = [ "compress=zstd" "noatime" ];
                      };
                      "@/var-lib" = {
                        mountpoint = "/var/lib";
                        mountOptions = [ "compress=zstd" "noatime" ];
                      };
                      "@/var-log" = {
                        mountpoint = "/var/log";
                        mountOptions = [ "compress=zstd" "noatime" ];
                      };
                      "@/var-tmp" = {
                        mountpoint = "/var/tmp";
                        mountOptions = [ "compress=zstd" "noatime" ];
                      };
                      "@swap" = lib.mkIf withSwap {
                        mountpoint = "/.swapvol";
                        swap.swapfile.size = "${swapSize}G";
                      };
                    }
                  else
                    {
                      "@" = { };
                      "@/root" = {
                        mountpoint = "/";
                        mountOptions = [ "compress=zstd" "noatime" ];
                      };
                      "@/nix" = {
                        mountpoint = "/nix";
                        mountOptions = [ "compress=zstd" "noatime" ];
                      };
                      "@/persist" = {
                        mountpoint = "/persist";
                        mountOptions = [ "compress=zstd" "noatime" ];
                      };
                      "@/var-lib" = {
                        mountpoint = "/var/lib";
                        mountOptions = [ "compress=zstd" "noatime" ];
                      };
                      "@/var-log" = {
                        mountpoint = "/var/log";
                        mountOptions = [ "compress=zstd" "noatime" ];
                      };
                      "@/var-tmp" = {
                        mountpoint = "/var/tmp";
                        mountOptions = [ "compress=zstd" "noatime" ];
                      };
                      # TODO: Check if this is correct implementation
                      "@swap" = lib.mkIf withSwap {
                        mountpoint = "/.swapvol";
                        swap.swapfile.size = "${swapSize}G";
                      };
                    };
              };
            };
          };
        };
      };

      vdb = if (number_of_disks == 1) then {}
      else
      {
        type = "disk";
        device = builtins.elemAt disks 1;
        content = {
          type = "gpt";
          partitions = {
            DATA = {
              # name = "DATA";
              start = "1MiB";
              end = "100%";
              content = {
                type = "btrfs";
                extraArgs = [ "-f" ]; # Override existing partition
                subvolumes = {
                  "@" = { 
                    mountpoint = "/DATA";
                    mountOptions = [ "compress=zstd" "noatime" ];
                  };
                  "@/home" = {
                    mountpoint = "/home";
                    mountOptions = [ "compress=zstd" ];
                  };
                };
              };
            };
          };
        };
      };
    };
  };
}