# 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)
#   bootScript  — sh script run as root on the BOOTED image through Apple's QEMU
#                 guest agent (network, user session available; run after `script`)
#   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
#   network     — attach a user-mode NIC for bootScript (default true)
{ pkgs, lib, qemu, ident, makeVmixVolume, makeOpenCore, makeBootDisk, installBootloader, vmixReadback, vmDriver, ... }:
originalImage: {
  name ? "",
  script ? "",
  bootScript ? "",
  network ? true,
  files ? [],
  smbios ? null,
  diskSize ? "",
  impure ? true,
  vncDisplay ? null,
  smp ? 4,
  memSize ? 4096,
  cpu ? qemu.defaultCpu,
  timeout ? 1800,
  machineArgs ? null,   # override qemu.machineArgs (device experiments)
}:
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 != "";
  hasBootScript = bootScript != "";
  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;
  };
  bootRunScript = pkgs.writeText "${name}-boot.sh" ''
    #!/bin/bash
    # runs as root on the booted system (guest-exec); VMIX is mounted at $V
    V=/Volumes/VMIX
    echo "=== vmix (online): ${name} ==="
    CONSOLE_USER=$(stat -f %Su /dev/console 2>/dev/null)
    CONSOLE_UID=$(id -u "$CONSOLE_USER" 2>/dev/null)
    export V CONSOLE_USER CONSOLE_UID
    # run something inside the logged-in user's GUI session
    as_user() { launchctl asuser "$CONSOLE_UID" sudo -u "$CONSOLE_USER" "$@"; }
    ${bootScript}
  '';
  bootVol = makeVmixVolume {
    name = "${name}-${originalImageName}-boot";
    files = [ { source = bootRunScript; name = "run.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 \
      ${if machineArgs != null then machineArgs else 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 ==="
  '';

  onlineCommands = lib.optionalString hasBootScript ''
    cp ${bootVol} vmix-boot.img
    chmod +w vmix-boot.img
    cp ${pkgs.OVMF.fd}/FV/OVMF_VARS.fd vars-boot.fd
    chmod +w vars-boot.fd
    VMIX_DISPLAY="-display none"
    ${lib.optionalString (vncDisplay != null) ''VMIX_DISPLAY="-display none -vnc ${vncDisplay}"''}
    QGA_SOCK=$(mktemp -u /tmp/vmix-qga-XXXXXX.sock)

    echo "=== vmix: booting ${originalImageName} for ${name} (guest agent) ==="
    python3 ${vmDriver} --mode qga --name "${name}-${originalImageName}-online" --timeout ${toString timeout} \
      --serial-log serial-boot.log --qga-sock "$QGA_SOCK" \
      --qga-command 'for i in $(seq 1 30); do diskutil mount VMIX >/dev/null 2>&1; [ -f /Volumes/VMIX/run.sh ] && break; sleep 2; done; [ -f /Volumes/VMIX/run.sh ] || { echo "no VMIX volume"; exit 9; }; bash /Volumes/VMIX/run.sh > /Volumes/VMIX/vmix-run.log 2>&1; rc=$?; echo $rc > /Volumes/VMIX/vmix-run.status; sync; cat /Volumes/VMIX/vmix-run.log; diskutil unmount force /Volumes/VMIX >/dev/null 2>&1; exit $rc' -- \
      qemu-system-x86_64 $VMIX_DISPLAY \
      ${if machineArgs != null then machineArgs else qemu.machineArgs { inherit cpu smp memSize; }} \
      ${qemu.firmwareArgs "vars-boot.fd"} \
      ${qemu.serialArgs "serial-boot.log"} \
      ${qemu.guestAgentArgs "$QGA_SOCK"} \
      ${qemu.sataDrive { id = "system"; port = 0; file = resultImg; }} \
      ${qemu.sataDrive { id = "vmix"; port = 1; file = "vmix-boot.img"; format = "raw"; }} \
      ${lib.optionalString network (qemu.netArgs { mac = originalImage.macAddress; })} \
      || { echo "vmix: online step failed during ${name} (see /tmp/vmix-macos/${name}-${originalImageName}-online)"; exit 1; }
    rm -f "$QGA_SOCK"
    ${vmixReadback "vmix-boot.img"}
    [ "$STATUS" = "0" ] || { echo "vmix: ${name} bootScript failed (status '$STATUS')"; exit 1; }
    echo "=== vmix: ${name} (online) 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}
    ${onlineCommands}
    ${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; }