{ config, pkgs, lib, inputs, modulesPath, ... }:

# prior art:
# https://github.com/dramforever/nixos-riscv-efi/
# https://github.com/NixOS/nixos-hardware/blob/master/starfive/visionfive/v2/sd-image.nix
# https://github.com/NixOS/nixpkgs/blob/master/nixos/modules/installer/sd-card/sd-image.nix

# info:
# https://u-boot.readthedocs.io/en/latest/develop/spl.html
# https://github.com/starfive-tech/VisionFive2
# https://github.com/u-boot/u-boot/blob/master/doc/board/starfive/visionfive2.rst
# https://github.com/starfive-tech/u-boot/commit/2a040b7720825629103c5b1efc5ace023072db8d
# https://www.man7.org/linux/man-pages/man8/sfdisk.8.html

# guidance
# https://reddit.com/r/RISCV/comments/zy30vl/comment/j3zhzkx/
# https://forum.rvspace.org/t/nvme-boot-using-visionfive2-software-v2-11-5/2464/1
# https://forum.rvspace.org/t/i-tried-the-v3-0-4-sdcard-img-from-sdcard-and-from-nvme/2965


let
  firmware = pkgs.callPackage "${inputs.nixos-hardware}/starfive/visionfive/v2/firmware.nix" { };
in
{
  options.efiImage = with lib; {

    imageBaseName = mkOption {
      default = "nixos-efi-image";
      description = mdDoc "Prefix of the name of the generated image file.";
    };

    compressImage = mkOption {
      type = types.bool;
      default = false;
      description = mdDoc "Whether the image should be compressed using {command}`zstd`.";
    };

  };

  config = {

    #assertions = [
    #  {
    #    assertion = config.boot.loader.systemd-boot.enable;
    #    message = "Building EFI Image requires systemd-boot";
    #  }
    #];

    #fileSystems.${config.boot.loader.efi.efiSysMountPoint}.neededForBoot = true;
    #fileSystems."/".autoResize = true;
    #boot.growPartition = true;

    system.build.efiImage = pkgs.runCommand "efi-image" {
      nativeBuildInputs = with pkgs; [
        util-linux # (partx)
      ] ++ lib.optional config.efiImage.compressImage zstd;
      rootImage = pkgs.callPackage "${modulesPath}/../lib/make-ext4-fs.nix" {
        storePaths = [ config.system.build.toplevel ];
        volumeLabel = "NIXOS_ROOT";
        populateImageCommands = ''
          mkdir -p ./files/boot
          ${config.boot.loader.generic-extlinux-compatible.populateCmd} -c ${config.system.build.toplevel} -d ./files/boot
        '';
      };
      skipSize = 1;   # initial offset
      espSize  = 512; # reserved boot partition size
      imageName = "${config.efiImage.imageBaseName}-${config.system.nixos.label}-${pkgs.stdenv.hostPlatform.system}-starfive-visionfive2.img";
    } ''
      mkdir -p "$out"
      img="$out/$imageName"

      rootSizeBlocks="$(du -B 512 --apparent-size "$rootImage" | awk '{ print $1 }')"
      espSizeBlocks="$(( espSize * 1024 * 1024 / 512 ))"
      imageSize=$(( rootSizeBlocks * 512 + espSizeBlocks * 512 + skipSize * 1024 * 1024 + 1024 * 1024 ))

      truncate -s "$imageSize" "$img"

      sfdisk $img <<EOF
        label: gpt
        unit: sectors
        sector-size: 512

        start=4096, size=4096, type=2E54B353-1271-4842-806F-E436D6AF6985, name=UBOOT_SPL
        start=8192, size=8192, type=5B193300-FC78-40CD-8002-E86C45580B47, name=UBOOT_FIRMWARE
        start=$(( skipSize + espSize ))M, type=linux, attrs="LegacyBIOSBootable"
      EOF

      eval "$(partx $img -o START,SECTORS --nr 1 --pairs)"
      dd conv=notrunc if=${firmware.spl}/share/starfive-visionfive2/spl.bin of=$img seek=$START count=$SECTORS

      eval "$(partx $img -o START,SECTORS --nr 2 --pairs)"
      dd conv=notrunc if=${firmware.uboot-fit-image}/share/starfive-visionfive2/visionfive2_fw_payload.img of=$img seek=$START count=$SECTORS

      eval "$(partx $img -o START,SECTORS --nr 3 --pairs)"
      dd conv=notrunc if="$rootImage" of="$img" seek="$START" count="$SECTORS"

      if test -n "$compressImage"; then
        zstd -T$NIX_BUILD_CORES --rm $img
      fi
    '';

    boot.postBootCommands = ''
      # On the first boot do some maintenance tasks
      if [ -f /nix-path-registration ]; then
        set -euo pipefail
        set -x

        # Register the contents of the initial Nix store
        ${config.nix.package.out}/bin/nix-store --load-db < /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

        # create the first boot profile
        NIXOS_INSTALL_BOOTLOADER=1 /nix/var/nix/profiles/system/bin/switch-to-configuration boot
        rm -f /boot/loader/entries/initial-nixos.conf

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

  };

}