# Customize Windows images via offline registry merge and/or Audit Mode script execution.
#
# Templates can provide:
#   windowsRegistry — merged offline via virt-win-reg (fast, no boot needed)
#   auditScript     — injected into image and run in Audit Mode via QEMU boot
#   cdroms          — ISO files to attach as CDs when booting for auditScript
#
# Both can be combined: registry is applied first (offline), then the script runs (online).
{ pkgs, lib, ... }:
  originalImage: {
    name ? "",
    diskSize ? "",
    impure ? true,
    # Offline: merge .reg file into registry via virt-win-reg
    windowsRegistry ? "",
    # Online: boot into Audit Mode and run this script
    auditScript ? "",
    # CD-ROMs to attach when booting for auditScript (e.g. VirtIO ISO)
    cdroms ? [],
    # Files to upload into the image before running auditScript
    # List of { source = <drv or path>; dest = "/Windows/path"; }
    uploads ? [],
    # QEMU settings for auditScript boot
    vncDisplay ? null,
    smp ? 4,
    memSize ? 4096,
  }:
  let
    originalImageName = lib.strings.removeSuffix "-vmix" (lib.strings.removeSuffix ".qcow2" originalImage.name);
    customImageName = (if name != "" then name else "custom") + "-${originalImageName}-vmix.qcow2";
    resultImg = "./disk.qcow2";

    isAHCI = originalImage.useAHCI or false;
    hasRegistry = windowsRegistry != "";
    hasAuditScript = auditScript != "";

    # Offline registry merge
    windowsRegFile = pkgs.writeText "${name}-registry.reg" windowsRegistry;
    virtWinRegMerge = lib.optionalString hasRegistry ''

      echo "=== vmix: merging registry entries (${name}) ==="
      virt-win-reg --merge ${resultImg} ${windowsRegFile}
    '';

    # Audit Mode script injection + QEMU boot
    auditScriptFile = pkgs.writeText "${name}-audit.cmd" auditScript;
    wrapperScript = pkgs.writeText "${name}-audit-wrapper.cmd" ''
      @echo off
      echo === vmix audit: ${name} ===
      call C:\vmix-audit-script.cmd
      echo === vmix audit: ${name} complete ===
      del /q C:\vmix-audit-script.cmd 2>nul
      shutdown /s /t 5 /c "vmix: ${name} complete" 2>nul
      del /q C:\vmix-audit-wrapper.cmd 2>nul
    '';
    runOnceReg = pkgs.writeText "${name}-runonce.reg" (lib.concatStringsSep "\n" [
      "Windows Registry Editor Version 5.00"
      ""
      ''[HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\Windows\CurrentVersion\RunOnce]''
      ''"vmixAudit"="C:\\vmix-audit-wrapper.cmd"''
      ""
    ]);

    cdromArgs = lib.concatMapStringsSep " \\\n      " (cd: "-drive file=${cd},media=cdrom,readonly=on") cdroms;

    displayArg = if vncDisplay != null then "-vnc ${vncDisplay}" else null;

    auditBootCommands = lib.optionalString hasAuditScript ''

      echo "=== vmix: injecting audit script (${name}) ==="
      virt-customize -a ${resultImg} \
        --upload ${auditScriptFile}:/vmix-audit-script.cmd \
        --upload ${wrapperScript}:/vmix-audit-wrapper.cmd \
        ${lib.concatMapStringsSep " \\\n        " (u: let dir = builtins.dirOf u.dest; in "${lib.optionalString (dir != "/") "--mkdir ${dir}"} --upload ${u.source}:${u.dest}") uploads}

      echo "=== vmix: adding RunOnce entry ==="
      virt-win-reg --merge ${resultImg} ${runOnceReg}

      # Boot into Audit Mode to run the script
      cp ${pkgs.OVMF.fd}/FV/OVMF_VARS.fd vars.fd
      chmod +w vars.fd

      VMIX_DISPLAY="-nographic"
      ${lib.optionalString (vncDisplay != null) ''VMIX_DISPLAY="-vnc ${vncDisplay}"''}
      ${lib.optionalString (vncDisplay == null) ''
        VMIX_DF=$(ls -t /tmp/.vmix-display-* 2>/dev/null | head -1)
        if [ -n "$VMIX_DF" ]; then
          export DISPLAY=$(sed -n '1p' "$VMIX_DF")
          export XAUTHORITY=$(sed -n '2p' "$VMIX_DF")
          export HOME=$(mktemp -d)
          export XDG_RUNTIME_DIR=$HOME
          export SDL_VIDEODRIVER=x11
          VMIX_DISPLAY="-display sdl"
        fi
      ''}

      echo "=== vmix: booting Audit Mode for ${name} ==="
      QEMU_ARGS="-accel kvm -m ${toString memSize} -smp ${toString smp} -cpu host -machine type=q35 \
        -drive if=pflash,format=raw,readonly=on,file=${pkgs.OVMF.fd}/FV/OVMF_CODE.fd \
        -drive if=pflash,format=raw,file=vars.fd \
        -rtc base=localtime,clock=host \
        -device qemu-xhci -device usb-tablet \
        ${if isAHCI
          then "-drive file=${resultImg},format=qcow2,if=none,id=disk0 -device ide-hd,drive=disk0"
          else "-drive file=${resultImg},format=qcow2,if=virtio"} \
        ${cdromArgs} \
        -nic user,model=${if isAHCI then "e1000" else "virtio-net-pci"}"

      timeout 1800 qemu-system-x86_64 $VMIX_DISPLAY $QEMU_ARGS || \
        if [[ "$VMIX_DISPLAY" == "-display sdl" ]]; then
          echo "=== vmix: SDL failed, retrying headless ==="
          cp ${pkgs.OVMF.fd}/FV/OVMF_VARS.fd vars.fd && chmod +w vars.fd
          timeout 1800 qemu-system-x86_64 -nographic $QEMU_ARGS
        else
          exit 1
        fi

      echo "=== vmix: audit script ${name} complete ==="
    '';

    builderCommand = ''
      # create resulting image backed by original image
      qemu-img create -f qcow2 -b ${originalImage} -F qcow2 ${resultImg}
      [ -n "${diskSize}" ] && qemu-img resize ${resultImg} ${diskSize}
      ${virtWinRegMerge}
      ${auditBootCommands}
      mv ${resultImg} $out
    '';
  builtImage = pkgs.runCommand customImageName ({
    nativeBuildInputs = with pkgs; [ qemu perl guestfs-tools ];
    requiredSystemFeatures = [ "kvm" ];
  } // lib.optionalAttrs impure { __noChroot = true; }) builderCommand;
  in
    builtImage // { _vmixOsType = "windows"; useAHCI = isAHCI; }