# Based on `nixos/modules/installer/sd-card/sd-image.nix`
{ lib, modulesPath, pkgs, config, ... }:
let
  inherit (lib) concatMap mkOption;

  cfg = config.my.disk;

  nixImage = (pkgs.callPackage "${modulesPath}/../lib/make-ext4-fs.nix" {
    volumeLabel = "qclkos-nix";
    uuid = "38bd3706-c049-430d-9e33-5cd0e437279b";
    storePaths = config.system.build.toplevel;
    compressImage = false;
  }).overrideAttrs (o: {
    buildCommand = ''
      # We need to move the store up a level since we're not using this as a rootfs but as /nix
      # HACK: `populateImageCommands` is executed in a subshell _before_ the store paths are copied in...
      shopt -s expand_aliases
      nixUpThenFaketime() {
        mv rootImage/{nix/store,}
        rmdir rootImage/nix

        faketime "$@"
      }
      alias faketime=nixUpThenFaketime

      ${o.buildCommand}
    '';
  });
in
{
  options.my.disk = with lib.types; {
    bootSize = mkOption {
      description = "/boot size (MiB).";
      type = ints.unsigned;
      default = 1024;
    };
    persistSize = mkOption {
      description = "/persist size (MiB).";
      type = ints.unsigned;
      default = 8192;
    };

    populateBootCommands = mkOption {
      description = ''
        Shell commands to populate the ./boot directory.
        All files in that directory are copied to the
        /boot partition on the image.
      '';
    };

    imageBaseName = mkOption {
      description = "Prefix of the name of the generated image file.";
      default = "qclkos";
    };
    image = mkOption {
      description = "Output disk image.";
      type = unspecified;
    };

    rootSize = mkOption {
      description = "tmpfs root size.";
      type = str;
      default = "2G";
    };
  };

  config = {
    fileSystems = {
      "/boot" = {
        device = "/dev/disk/by-label/QCLKOS_BOOT";
        fsType = "vfat";
      };
      "/persist" = {
        device = "/dev/disk/by-label/qclkos-persist";
        fsType = "ext4";
        neededForBoot = true;
      };
      "/nix" = {
        device = "/dev/disk/by-label/qclkos-nix";
        fsType = "ext4";
      };

      "/" = {
        device = "yeet";
        fsType = "tmpfs";
        options = [ "size=${cfg.rootSize}" "mode=755" ];
      };
    };

    boot = {
      postBootCommands = ''
        # On the first boot do some maintenance tasks
        if [ -f /nix/nix-path-registration ]; then
          set -euo pipefail
          set -x
          # Figure out device names for the boot device and nix filesystem.
          nixPart=$(${pkgs.util-linux}/bin/findmnt -n -o SOURCE /nix)
          bootDevice=$(lsblk -npo PKNAME $nixPart)
          partNum=$(lsblk -npo MAJ:MIN $nixPart | ${pkgs.gawk}/bin/awk -F: '{print $2}')

          # Resize the nix partition and the filesystem to fit the disk
          echo ",+," | sfdisk -N$partNum --no-reread $bootDevice
          ${pkgs.parted}/bin/partprobe
          ${pkgs.e2fsprogs}/bin/resize2fs $nixPart

          # Register the contents of the initial Nix store
          ${config.nix.package.out}/bin/nix-store --load-db < /nix/nix-path-registration

          # nixos-rebuild also requires a "system" profile and an /etc/NIXOS tag.
          touch /etc/NIXOS
          ${config.nix.package.out}/bin/nix-env -p /nix/var/nix/profiles/system --set /run/current-system

          # Prevents this from running on later boots.
          rm -f /nix/nix-path-registration
        fi
      '';
    };

    users.mutableUsers = false;

    environment = {
      # This is based on parts of `nixfiles/nixos/modules/tmproot.nix`
      persistence."/persist" = {
        directories = [
          "/var/lib/nixos"
          "/var/lib/systemd"

          "/etc/qclk"
        ];
        files = [
          "/etc/machine-id"
        ] ++ (concatMap (k: [ k.path "${k.path}.pub" ]) config.services.openssh.hostKeys);
      };
    };

    services = {
      journald.storage = "volatile";
    };

    my.disk.image = pkgs.callPackage ({
      stdenv, dosfstools, e2fsprogs, mtools, libfaketime, util-linux
    }: stdenv.mkDerivation {
      name = "${cfg.imageBaseName}-${config.system.nixos.label}-${pkgs.stdenv.hostPlatform.system}";
      nativeBuildInputs = [
        dosfstools e2fsprogs libfaketime mtools util-linux
      ];

      buildCommand = ''
        mkdir -p $out
        export img=$out/${cfg.imageBaseName}.img

        nix_fs=${nixImage}

        # Gap in front of the first partition, in MiB
        gap=8

        # Create the image file sized to fit /boot and /nix, plus slack for the gap.
        nixSizeBlocks=$(du -B 512 --apparent-size $nix_fs | awk '{ print $1 }')
        bootSizeBlocks=$((${toString cfg.bootSize} * 1024 * 1024 / 512))
        persistSizeBlocks=$((${toString cfg.persistSize} * 1024 * 1024 / 512))
        imageSize=$((nixSizeBlocks * 512 + persistSizeBlocks * 512 + bootSizeBlocks * 512 + gap * 1024 * 1024))
        truncate -s $imageSize $img

        # type=b is 'W95 FAT32', type=83 is 'Linux'.
        # The "bootable" partition is where u-boot will look file for the bootloader
        # information (dtbs, extlinux.conf file).
        sfdisk --no-reread --no-tell-kernel $img <<EOF
          label: dos
          label-id: 0xabb10d5a

          start=''${gap}M, size=$bootSizeBlocks, type=b, bootable
          start=$((gap + ${toString cfg.bootSize}))M, size=$persistSizeBlocks, type=83
          start=$((gap + ${toString cfg.bootSize} + ${toString cfg.persistSize}))M, type=83
        EOF

        # Copy the nixfs into the image
        eval $(partx $img -o START,SECTORS --nr 3 --pairs)
        dd conv=notrunc if=$nix_fs of=$img seek=$START count=$SECTORS

        # Create a FAT32 /boot partition of suitable size into boot_part.img
        eval $(partx $img -o START,SECTORS --nr 1 --pairs)
        truncate -s $((SECTORS * 512)) boot_part.img

        mkfs.vfat --invariant -n QCLKOS_BOOT boot_part.img

        # Populate the files intended for /boot
        mkdir boot
        ${cfg.populateBootCommands}

        find boot -exec touch --date=2000-01-01 {} +
        # Copy the populated /boot into the image
        cd boot
        # Force a fixed order in mcopy for better determinism, and avoid file globbing
        for d in $(find . -type d -mindepth 1 | sort); do
          faketime "2000-01-01 00:00:00" mmd -i ../boot_part.img "::/$d"
        done
        for f in $(find . -type f | sort); do
          mcopy -pvm -i ../boot_part.img "$f" "::/$f"
        done
        cd ..

        # Verify the FAT partition before copying it.
        fsck.vfat -vn boot_part.img
        dd conv=notrunc if=boot_part.img of=$img seek=$START count=$SECTORS
      '';
    }) { };
  };
}