# Customize a macOS image offline from the vmix PE: the recovery boots with the
# image and a VMIX volume attached, its hook runs `script` as root with the
# image's System (read-only) and Data (rw) volumes mounted at $SYS / $DATA, then
# powers off. The installed macOS itself is never booted, so nothing depends on
# launchd/BTM approval inside the guest. Counterpart of the Windows
# registry/audit flow. Optionally re-installs OpenCore with a new SMBIOS
# identity (`smbios`).
#
# Templates provide:
#   script  — sh script run as root in the PE (pe-lib.sh helpers available)
#   files   — [{ source; name; }] extra files placed next to it on /Volumes/VMIX
#   smbios  — { model? serial? mlb? uuid? mac? seed? } → fresh OpenCore config in the ESP
{ pkgs, lib, qemu, ident, makeVmixVolume, makeOpenCore, makeBootDisk, installBootloader, vmixReadback, vmDriver, ... }:
originalImage: {
  name ? "",
  script ? "",
  files ? [],
  smbios ? null,
  diskSize ? "",
  impure ? true,
  vncDisplay ? null,
  smp ? 4,
  memSize ? 4096,
  cpu ? qemu.defaultCpu,
  timeout ? 1800,
}:
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";
  hasScript = script != "";
  hasSmbios = smbios != null;
  pe = originalImage.pe or (throw "vmix: image ${originalImage.name} carries no PE (built by an older makeImage?)");
  volumeName = originalImage.volumeName or "Macintosh HD";

  model = originalImage.model or "MacPro7,1";
  seed = if hasSmbios && (smbios.seed or null) != null then smbios.seed else null;
  mac = if !hasSmbios then originalImage.macAddress
        else if (smbios.mac or null) != null then smbios.mac
        else if seed != null then ident.macFromSeed seed
        else originalImage.macAddress;
  uuid = if !hasSmbios then null
         else if (smbios.uuid or null) != null then smbios.uuid
         else if seed != null then ident.uuidFromSeed seed
         else originalImage.opencore.uuid;
  esp = if hasSmbios
    then makeOpenCore ({
        name = "${name}-${originalImageName}-opencore";
        model = smbios.model or model;
        inherit mac uuid;
      } // builtins.removeAttrs smbios [ "seed" "mac" "uuid" "model" ])
    else originalImage.opencore;
  # PE boot disk: serial console, and an OpenCore ScanPolicy that only allows
  # HFS+ volumes on SATA (= the PE), so the image's own macOS is never booted.
  # 0x10203 = FILE_SYSTEM_LOCK | DEVICE_LOCK | ALLOW_FS_HFS | ALLOW_DEVICE_SATA
  bootDisk = makeBootDisk {
    name = "${name}-${originalImageName}-pe";
    esp = originalImage.opencore;
    bootArgs = "keepsyms=1 serial=3 -v";
    scanPolicy = 66051;
  };

  runScript = pkgs.writeText "${name}-run.sh" ''
    #!/bin/bash
    . /Volumes/VMIX/pe-lib.sh
    echo "=== vmix: ${name} ==="
    pe_mount_target || pe_fail "could not mount the target volumes"
    ${script}
    pe_unmount_target
  '';
  vmixVol = makeVmixVolume {
    name = "${name}-${originalImageName}";
    files = [
      { source = runScript; name = "run.sh"; }
      { source = ../guest/pe-lib.sh; name = "pe-lib.sh"; }
    ] ++ files;
  };
  driverPython = pkgs.python3.withPackages (p: [ p.pillow ]);

  bootCommands = lib.optionalString hasScript ''
    cp ${vmixVol} vmix.img
    chmod +w vmix.img
    cat > vmix.conf <<CONF
    VOLUME_NAME="${volumeName}"
    BUILD_DATE="$(date -u +%m%d%H%M%Y.%S)"
    CONF
    guestfish -a vmix.img -m /dev/sda1 upload vmix.conf /vmix.conf
    qemu-img create -q -f qcow2 -F raw -b ${pe} pe.qcow2
    qemu-img create -q -f qcow2 -F raw -b ${bootDisk}/boot.img ocboot.qcow2
    cp ${pkgs.OVMF.fd}/FV/OVMF_VARS.fd vars.fd
    chmod +w vars.fd
    VMIX_DISPLAY="-display none"
    ${lib.optionalString (vncDisplay != null) ''VMIX_DISPLAY="-display none -vnc ${vncDisplay}"''}

    echo "=== vmix: running ${name} in the PE against ${originalImageName} ==="
    python3 ${vmDriver} --mode pe --name "${name}-${originalImageName}" --timeout ${toString timeout} \
      --serial-log serial.log --progress-file ${resultImg} -- \
      qemu-system-x86_64 $VMIX_DISPLAY \
      ${qemu.machineArgs { inherit cpu smp memSize; }} \
      ${qemu.firmwareArgs "vars.fd"} \
      ${qemu.serialArgs "serial.log"} \
      ${qemu.sataDrive { id = "opencore"; port = 0; file = "ocboot.qcow2"; }} \
      ${qemu.sataDrive { id = "pe"; port = 1; file = "pe.qcow2"; }} \
      ${qemu.sataDrive { id = "system"; port = 2; file = resultImg; }} \
      ${qemu.sataDrive { id = "vmix"; port = 3; file = "vmix.img"; format = "raw"; }} \
      || { echo "vmix: PE failed during ${name} (see /tmp/vmix-macos/${name}-${originalImageName})"; exit 1; }

    ${vmixReadback "vmix.img"}
    [ "$STATUS" = "0" ] || { echo "vmix: ${name} script failed (status '$STATUS')"; exit 1; }
    echo "=== vmix: ${name} complete ==="
  '';

  builtImage = pkgs.runCommand customImageName ({
    nativeBuildInputs = with pkgs; [ pkgs.qemu driverPython libguestfs-with-appliance ];
    requiredSystemFeatures = [ "kvm" ];
  } // lib.optionalAttrs impure { __noChroot = true; }) ''
    qemu-img create -q -f qcow2 -b ${originalImage} -F qcow2 ${resultImg}
    [ -n "${diskSize}" ] && qemu-img resize ${resultImg} ${diskSize}
    ${bootCommands}
    ${lib.optionalString hasSmbios (installBootloader { inherit esp; image = resultImg; })}
    mv ${resultImg} $out
  '';
in
  builtImage // { _vmixOsType = "macos"; macAddress = mac; opencore = esp; model = esp.model or model; inherit pe volumeName; }