# Local copy of the clan-core module in patches/upstream-sd-card-secret-injection.patch.
# Delete once clan-core ships clan.core.image.sd-card.addFilesScript
# (the option declaration would conflict).
{
  inputs,
  lib,
  pkgs,
  ...
}:
let
  parseAddFilesArgs = "${inputs.clan-core}/nixosModules/clanCore/image/parse-add-files-args.sh";

  # Writes extra files straight into the ext4 root partition of an SD image
  # with debugfs, so no boot-time hook is needed: the files exist on the root
  # filesystem from the first boot. Runs unprivileged and without mounting.
  sdCardAddFilesScript = pkgs.writeShellApplication {
    name = "add-files";
    runtimeInputs = with pkgs; [
      coreutils
      diffutils
      e2fsprogs
      findutils
      util-linux
      zstd
    ];
    extraShellCheckFlags = [ "-x" ];
    text = ''
      # shellcheck source=${parseAddFilesArgs}
      source ${parseAddFilesArgs}

      # If no extra paths, just copy source to output
      if [[ ''${#extra_paths[@]} -eq 0 ]]; then
        cp -- "$source_image" "$output_image"
        exit 0
      fi

      # The injection sandbox mounts /tmp as a RAM-backed tmpfs and a raw SD
      # image is several GiB, so stage it next to the output, the only
      # writable on-disk path.
      workdir=$(dirname "$output_image")
      img="$workdir/.$(basename "$output_image").raw"
      cmds=$(mktemp)
      trap 'rm -f -- "$img" "$cmds"' EXIT

      case "$source_image" in
        *.zst)
          zstd -d -q --no-progress -o "$img" -- "$source_image"
          compress=1
          ;;
        *)
          cp --no-preserve=mode -- "$source_image" "$img"
          compress=0
          ;;
      esac
      chmod u+w -- "$img"

      # sd-image.nix writes an MBR whose ext4 root is the (only) type-83 partition.
      root_start=""
      while read -r _ start type; do
        if [[ "$type" == "0x83" ]]; then
          root_start=$start
        fi
      done < <(partx --show -g -o NR,START,TYPE -- "$img")
      if [[ -z "$root_start" ]]; then
        echo "No Linux (type 83) root partition found in $source_image" >&2
        exit 1
      fi
      # e2fsprogs opens a partition inside a disk image via "<file>?offset=<bytes>".
      rootfs="$img?offset=$((root_start * 512))"

      # Build a debugfs command script. debugfs has neither mkdir -p nor a
      # recursive write, so directories are created ancestor-first ("already
      # exists" errors are harmless). write copies the source mode; ownership
      # is forced to root since the script runs as an unprivileged user.
      declare -A made_dirs=()
      written=()
      emit_mkdir_p() {
        local dir="$1" path="" part
        IFS=/ read -ra parts <<< "$dir"
        for part in "''${parts[@]}"; do
          [[ -z "$part" ]] && continue
          path="$path/$part"
          if [[ -z "''${made_dirs[$path]:-}" ]]; then
            echo "mkdir \"$path\""
            made_dirs[$path]=1
          fi
        done
      }
      emit_write() {
        local src="$1" dest="$2"
        echo "write \"$src\" \"$dest\""
        echo "sif \"$dest\" uid 0"
        echo "sif \"$dest\" gid 0"
        written+=("$src" "$dest")
      }

      for ((i = 0; i < ''${#extra_paths[@]}; i += 2)); do
        local_src="''${extra_paths[i]}"
        # Normalize to an absolute path without trailing slash; "/" becomes "".
        target="/''${extra_paths[i+1]#/}"
        target="''${target%/}"

        if [[ -d "$local_src" ]]; then
          emit_mkdir_p "$target"
          while IFS= read -r -d "" entry; do
            dest="$target/''${entry#"$local_src"/}"
            if [[ -d "$entry" ]]; then
              emit_mkdir_p "$dest"
            elif [[ -f "$entry" ]]; then
              emit_write "$entry" "$dest"
            else
              echo "Unsupported file type (only regular files and directories): $entry" >&2
              exit 1
            fi
          done < <(
            # Never treat our own staging/output directory as input, even if
            # the caller placed it inside an extra path.
            find "$local_src" -mindepth 1 -path "$workdir" -prune -o -print0
          )
        elif [[ -f "$local_src" ]]; then
          emit_mkdir_p "$(dirname "$target")"
          emit_write "$local_src" "$target"
        else
          echo "Extra path does not exist: $local_src" >&2
          exit 1
        fi
      done > "$cmds"

      debugfs -w -f "$cmds" "$rootfs"

      # debugfs exits 0 even when commands fail: verify every file landed
      # intact and the filesystem is still consistent.
      for ((i = 0; i < ''${#written[@]}; i += 2)); do
        if ! debugfs -R "cat \"''${written[i+1]}\"" "$rootfs" 2>/dev/null | cmp -s -- - "''${written[i]}"; then
          echo "Failed to write ''${written[i+1]} into the root filesystem" >&2
          exit 1
        fi
      done
      if ! fsck_output=$(EXT2FS_NO_MTAB_OK=yes e2fsck -fn "$rootfs" 2>&1); then
        echo "Root filesystem is inconsistent after injection:" >&2
        echo "$fsck_output" >&2
        exit 1
      fi

      if [[ "$compress" == 1 ]]; then
        zstd -T0 -q --no-progress --rm -o "$output_image" -- "$img"
      else
        mv -- "$img" "$output_image"
      fi
    '';
  };
in
{
  options.clan.core.image = {
    sd-card.addFilesScript = lib.mkOption {
      type = lib.types.package;
      default = sdCardAddFilesScript;
      defaultText = lib.literalExpression "pkgs.writeShellApplication { ... }";
      description = ''
        Script to inject extra files into an SD card image
        (`clan machines build <machine> --format sd-card`).

        Files are written directly into the ext4 root partition, so they are
        present from the first boot without any initrd hook. Both compressed
        (`.img.zst`, the `sdImage.compressImage` default) and raw `.img`
        images are supported; the output keeps the source's format.

        Source `parse-add-files-args.sh` to get standardized argument parsing.
        After sourcing, `$source_image`, `$output_image`, and `$extra_paths`
        are available.

        Accepts the following arguments:
        - --source `<path>`: Path to the original built image (read-only)
        - --output `<path>`: Path for the modified output image
        - --extra-path `<local-source>` `<target-path>`: Repeatable, maps a local
          file or directory to a target filesystem path inside the image
      '';
    };
  };
}