//! Per-vCPU worker and dispatch loop.
//!
//! The runner owns VM lifecycle. This module owns vCPU execution, MMIO exit
//! dispatch, PSCI handling, quiesce windows, and snapshot trigger checks.
//!
//! VM-execution module — gated to a platform with a hypervisor backend
//! (macOS/HVF today; broaden to include KVM when the Linux backend lands).
#![cfg(all(target_os = "macos", target_arch = "aarch64"))]

use std::fmt;
use std::sync::Arc;

use crate::devices::mmio_bus::MmioBus;
use crate::devices::virtio::mmio::MmioVirtio;
use crate::devices::virtio::vsock::device::Vsock;
use crate::fuse::DaxSession;
use crate::hypervisor::ActiveVcpu as Vcpu;
use crate::hypervisor::{CoreReg, HypervisorVcpu, SysReg, VcpuExit, VcpuHandle};
use crate::vmm::coord::VcpuCoordinator;
use crate::vmm::vstate::MicroVm;

pub struct DispatchSnapshot<'a> {
    pub after_ms: Option<u64>,
    pub at_heartbeat: Option<u64>,
    pub on_listener: bool,
    /// When `true`, the dispatch loop fires its readiness trigger
    /// on the per-VM
    /// [`crate::devices::serial::SerialState::pre_exec_ready`] atomic —
    /// init-oci just printed "workload-pre-exec" and is in a brief
    /// nanosleep, giving us a stable WFI window to capture in. Used
    /// by the always-pipelined-skip-warm `.build()` path: snapshots
    /// the guest BEFORE the workload runs, so each restore re-execs
    /// the workload fresh (which is what agent-runtime users want).
    /// Saves 50–150 ms of bake time vs `on_listener`.
    ///
    /// When `on_pre_exec` AND `on_listener` are both true, pre-exec
    /// wins (it fires earlier in the boot timeline). With `with_warmup`
    /// or service-image bakes, leave this `false` to preserve the
    /// existing listener-ready semantics.
    pub on_pre_exec: bool,
    pub quiesce_ms: u64,
    pub out_path: Option<&'a str>,
    pub stop_requested: Option<&'a std::sync::atomic::AtomicBool>,
    /// Pipelined-bake (`bake-then-pool`) signal. When `out_path`
    /// is `None` AND this is `true`, the dispatch loop returns
    /// `DispatchExit::BakeReady` on the first readiness trigger
    /// (workload-parked / listener-ready / snapshot-after-ms
    /// fallback) instead of doing the no-op `take_snapshot`. The
    /// runner's outer loop then writes `BAKE_READY` on the pool
    /// ctl writer and re-enters dispatch.
    pub bake_ready_signal: bool,
    /// Path to the worker's exec-vsock unix socket. When the
    /// snapshot trigger fires AND vCPUs > 1, the worker connects
    /// to this socket as a client, sends a `smpark_park` CONTROL
    /// (which the muxer forwards to the in-guest agent via
    /// virtio-vsock), waits for the ack, then proceeds with the
    /// pause-and-capture rendezvous. After the capture, sends
    /// `smpark_unpark` to wake the parked secondaries. Best-
    /// effort: `None` or transport failure or smpark.ko not loaded
    /// → fall through to the existing rendezvous-only path.
    pub vsock_exec_path: Option<&'a str>,
    /// This VM's per-VM PL011 state (`Vmm::serial`). The dispatch loop
    /// reads its boot-line markers (`heartbeat_count`, `pre_exec_ready`,
    /// `workload_parked`) as snapshot triggers and its
    /// `smpark_state_gpa` when writing the snapshot's smpark sidecar.
    pub serial: &'a std::sync::Arc<crate::devices::serial::SerialState>,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DispatchExit {
    Canceled,
    Stopped,
    SnapshotSaved,
    SystemOff,
    /// Returned by `dispatch_vcpu` exactly once per worker
    /// lifetime when `bake_ready_signal` is set: the bake init
    /// has reached its readiness trigger but we're not capturing
    /// (no `out_path`). Runner uses this to surface the signal
    /// to the host and transition to pool-mode dispatch.
    BakeReady,
}

#[derive(Debug)]
pub enum WorkerError {
    ListenerReadinessTimeout {
        after_ms: u64,
    },
    QuiesceThreadSpawn(std::io::Error),
    SnapshotCapture(crate::hypervisor::ActiveError),
    SnapshotSave {
        path: String,
        source: crate::vmm::snapshot::FileError,
    },
    VcpuCreate {
        idx: u32,
        source: crate::hypervisor::ActiveError,
    },
    VcpuRun {
        idx: u32,
        source: crate::hypervisor::ActiveError,
    },
    VcpuSetup {
        idx: u32,
        source: crate::hypervisor::ActiveError,
    },
    VtimerMask(crate::hypervisor::ActiveError),
}

impl fmt::Display for WorkerError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            WorkerError::ListenerReadinessTimeout { after_ms } => write!(
                f,
                "listener readiness timeout after {after_ms} ms; refusing zero-listener snapshot"
            ),
            WorkerError::QuiesceThreadSpawn(e) => write!(f, "spawn quiesce canceller: {e}"),
            WorkerError::SnapshotCapture(e) => write!(f, "capture snapshot: {e:?}"),
            WorkerError::SnapshotSave { path, source } => {
                write!(f, "save snapshot {path}: {source:?}")
            }
            WorkerError::VcpuCreate { idx, source } => {
                write!(f, "create vCPU {idx}: {source:?}")
            }
            WorkerError::VcpuRun { idx, source } => {
                write!(f, "run vCPU {idx}: {source:?}")
            }
            WorkerError::VcpuSetup { idx, source } => {
                write!(f, "setup vCPU {idx}: {source:?}")
            }
            WorkerError::VtimerMask(e) => write!(f, "set vtimer mask: {e:?}"),
        }
    }
}

impl std::error::Error for WorkerError {}

// The per-exit step outcome is the shared, portable `VcpuOutcome` (Phase 3 7b)
// — the vocabulary the KVM run loop will converge onto. Aliased locally so the
// existing `StepExit::*` call sites stay unchanged.
use crate::vcpu_dispatch::VcpuOutcome as StepExit;

/// Optionally hint macOS to schedule this vCPU thread onto a
/// P-core via `pthread_set_qos_class_self_np`. **Off by default**:
/// benchmarks (rust:1-slim warm rustc, 8 workers) showed the
/// scheduler's default behaviour beats both USER_INTERACTIVE and
/// USER_INITIATED for our workload — macOS already places HVF-
/// entitled foreground threads well, and forcing higher priority
/// only causes contention with other system threads. Kept as
/// opt-in for cases where the user is benchmarking on a quiet
/// host or has a specific workload that benefits.
///
/// Set `SUPERMACHINE_QOS=user_interactive` (or `=user_initiated`)
/// to enable. Anything else is a no-op.
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
pub fn pin_vcpu_thread_to_pcore() {
    const QOS_CLASS_USER_INTERACTIVE: u32 = 0x21;
    const QOS_CLASS_USER_INITIATED: u32 = 0x19;
    let cls = match std::env::var("SUPERMACHINE_QOS").as_deref() {
        Ok("user_interactive") => QOS_CLASS_USER_INTERACTIVE,
        Ok("user_initiated") => QOS_CLASS_USER_INITIATED,
        _ => return,
    };
    unsafe extern "C" {
        fn pthread_set_qos_class_self_np(qos_class: u32, relative_priority: i32) -> i32;
    }
    unsafe {
        let _ = pthread_set_qos_class_self_np(cls, 0);
    }
}

#[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
pub fn pin_vcpu_thread_to_pcore() {}

#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
pub fn run_secondary(
    idx: u32,
    coord: Arc<VcpuCoordinator>,
    bus: Arc<MmioBus>,
    dax_sessions: Vec<Arc<DaxSession>>,
    restore_state: Option<crate::vmm::snapshot::PerVcpuState>,
) {
    pin_vcpu_thread_to_pcore();
    if let Err(e) = run_secondary_inner(idx, coord, bus, dax_sessions, restore_state) {
        eprintln!("  [vcpu-{idx}] worker error: {e}");
    }
}

#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn run_secondary_inner(
    idx: u32,
    coord: Arc<VcpuCoordinator>,
    bus: Arc<MmioBus>,
    dax_sessions: Vec<Arc<DaxSession>>,
    restore_state: Option<crate::vmm::snapshot::PerVcpuState>,
) -> Result<(), WorkerError> {
    let vcpu = Vcpu::new().map_err(|source| WorkerError::VcpuCreate { idx, source })?;
    // MPIDR_EL1 must equal the vCPU index (matches FDT cpu@N reg + GIC
    // redistributor frame). Without this, GIC PPIs never route here.
    vcpu.set_sys(SysReg::MpidrEl1, idx as u64)
        .map_err(|source| WorkerError::VcpuSetup { idx, source })?;
    // Register the handle so the snapshot trigger thread can
    // force-exit us when it needs to capture state.
    coord.register_secondary(vcpu.exit_token());

    if let Some(st) = restore_state {
        // Multi-vCPU restore path. The snapshot already has our
        // full register state; load it and dispatch directly. No
        // PSCI park, no waiting for CPU_ON — the kernel onlined
        // us before snapshot fired and "thinks" we're still up,
        // which is true once we restore state and start running.
        //
        // BUT — first wait for the runner to finish
        // `hv_gic_set_state` on the boot thread. That call
        // restores the per-PE redistributor "global" state HVF
        // hides behind the opaque blob (WAKER ProcessorSleep,
        // CTLR enable, IGRPMODR0, NSACR, …). If we wrote our
        // per-PE GICR registers (IGROUPR0/ISENABLER0/ISPENDR0/…)
        // BEFORE the blob landed, the blob's later install would
        // overwrite our writes — or worse, on some HVF versions,
        // the blob's per-PE slots for not-yet-created vCPUs are
        // silently dropped, leaving this redistributor in
        // ProcessorSleep state. Either way the kernel sees a
        // redistributor that won't deliver IPIs, the vCPU can't
        // be woken from WFI, and any task scheduled on it stays
        // frozen. Manifested as intermittent (~13 % at 4 vCPUs,
        // 0 % at 1 vCPU) post-restore hangs in rosettad's UNIX
        // DGRAM recv when an amd64 binary tried to start.
        //
        // The runner's `cold_restore_with_secondaries` flow
        // spawns us BEFORE `hv_gic_set_state` so we exist as a
        // PE when the blob is applied; it then sets
        // `gic_restore_done` and we proceed.
        while !coord
            .gic_restore_done
            .load(std::sync::atomic::Ordering::Acquire)
        {
            if coord.shutdown.load(std::sync::atomic::Ordering::Acquire) {
                return Ok(());
            }
            std::hint::spin_loop();
        }
        eprintln!("  [vcpu-{idx}] restoring from snapshot");
        vcpu.restore_snapshot(&st)
            .map_err(|source| WorkerError::VcpuSetup { idx, source })?;
        // Align this secondary's CNTVOFF_EL2 with what the boot
        // vCPU's restore path published. Without this, the
        // secondary sees the host's raw CNTPCT (HVF default
        // CNTVOFF=0) while vcpu0 sees the snapshot-continued
        // CNTVCT — Linux's CLOCK_MONOTONIC then skews when a
        // userspace thread is migrated between cores, breaking
        // libuv's monotonic-time assertion in Node ≥24. coord
        // publishes the value as 0 before any restore happens;
        // we treat 0 as "no offset published, leave HVF default
        // alone" so cold-boot paths aren't disturbed.
        let off = coord
            .vtimer_offset
            .load(std::sync::atomic::Ordering::Acquire);
        if off != 0 {
            vcpu.set_vtimer_offset(off)
                .map_err(|source| WorkerError::VcpuSetup { idx, source })?;
        }
        coord.slots[idx as usize]
            .on
            .store(true, std::sync::atomic::Ordering::SeqCst);
        return dispatch_vcpu_inner(idx, &vcpu, &bus, &coord, &dax_sessions);
    }

    // Boot-from-scratch path. Secondary parks waiting for vcpu0
    // to issue PSCI CPU_ON during kernel SMP bring-up.
    eprintln!("  [vcpu-{idx}] parked, waiting for PSCI CPU_ON");
    let Some((entry, ctx_id)) = coord.wait_for_run(idx) else {
        return Ok(());
    };
    eprintln!("  [vcpu-{idx}] CPU_ON entry=0x{entry:x} ctx=0x{ctx_id:x}");
    // Linux aarch64 boot protocol for secondaries: same as primary:
    // EL1h, DAIF masked, MMU off. X0 = context id. X1..X3 = 0.
    vcpu.set_core(CoreReg::Cpsr, 0x3c5).ok();
    vcpu.set_core(CoreReg::Pc, entry).ok();
    vcpu.set_core(CoreReg::X(0), ctx_id).ok();
    vcpu.set_core(CoreReg::X(1), 0).ok();
    vcpu.set_core(CoreReg::X(2), 0).ok();
    vcpu.set_core(CoreReg::X(3), 0).ok();
    dispatch_vcpu_inner(idx, &vcpu, &bus, &coord, &dax_sessions)
}

/// Per-vCPU dispatch loop. Returns when the vCPU is canceled / shut down.
/// Snapshot trigger only runs on vCPU 0.
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
pub fn dispatch_vcpu(
    idx: u32,
    vcpu: &Vcpu,
    bus: &Arc<MmioBus>,
    coord: &Arc<VcpuCoordinator>,
    all_mmio: &[Arc<MmioVirtio>],
    posix_fs: &[Arc<crate::fuse::PosixFs>],
    dax_sessions: &[Arc<DaxSession>],
    fs_devices: &[Arc<crate::devices::virtio::fs::VirtioFs>],
    vsock: &Arc<Vsock>,
    vm: &MicroVm,
    snapshot: DispatchSnapshot<'_>,
) -> Result<DispatchExit, WorkerError> {
    use crate::vmm::snapshot;
    use std::sync::atomic::Ordering;

    let dispatch_start = std::time::Instant::now();
    let mut quiesced = false;
    let take_snapshot = |reason: &str| -> Result<bool, WorkerError> {
        let Some(out) = snapshot.out_path else {
            return Ok(false);
        };
        eprintln!("  snapshot trigger ({reason})");
        let t0 = std::time::Instant::now();
        // Multi-vCPU: park secondaries via smpark.ko before the
        // rendezvous. The worker connects to its own exec-vsock
        // unix socket; the muxer relays the CONTROL frame through
        // virtio-vsock to the in-guest agent, which ioctls
        // /dev/smpark. Secondaries broadcast-IPI into WFI with
        // IRQs masked + LRs drained, producing byte-identical-
        // trivial captured per-vCPU state that HVF can round-trip.
        // Best-effort: smpark.ko not loaded / no exec path / RPC
        // failure → fall through to existing rendezvous-only
        // capture (works on 1-vCPU; intermittent on N>1).
        // Note: smpark park/unpark is NOT driven from this in-
        // worker snapshot trigger path. The worker's vCPU 0
        // dispatcher thread is what would issue the agent RPC,
        // but it's also what runs hv_vcpu_run for vCPU 0 — so
        // the agent (which lives on vCPU 0) can never service
        // the request while we're blocking here. Result: 5s
        // timeout per snapshot. Multi-vCPU snapshot reliability
        // requires the host-driven pipelined-bake path; the
        // sequential bake's heartbeat trigger remains
        // intermittent-on-N>1.
        //
        // See bake.rs::run_native_supermachine_bake_pipelined
        // for the working path; PooledVm::snapshot also drives
        // smpark host-side from the embedder-facing API.
        let secondary_handles = coord.secondary_handles_snapshot();
        if !secondary_handles.is_empty() {
            coord.request_snapshot_pause(&secondary_handles);
        }
        // Park the virtio-fs I/O workers for the capture+save window:
        // they write FUSE replies into guest RAM from non-vCPU threads,
        // so the vCPU rendezvous alone no longer quiesces guest RAM.
        // Workers park at a FUSE-chain boundary; dropped when this
        // closure returns (after the streaming save + sidecars).
        let _fs_pause = crate::devices::virtio::fs::FsIoPauseGuard::new(fs_devices);
        let virtio = snapshot::VirtioSnapshot {
            // Unified device records (kind from device id, backing None — HVF
            // positional restore). 7c step 3.
            devices: all_mmio
                .iter()
                .map(|m| crate::snapshot_frame::DeviceRecord {
                    kind: crate::snapshot_frame::DeviceKind::from_virtio_id(m.device_id()),
                    mmio: m.capture_state(),
                    backing: crate::snapshot_frame::DeviceBacking::None,
                })
                .collect(),
            vsock_listeners: vsock.muxer().capture_tsi_listeners(),
        };
        let secondary_states = if !secondary_handles.is_empty() {
            coord.take_secondary_states()
        } else {
            Vec::new()
        };
        // Streaming capture+save: skips the 2 GiB intermediate
        // RAM copy that the legacy capture_snapshot allocates.
        // Same correctness invariant — guest is paused for the
        // duration via the secondary rendezvous above.
        let write_stats = snapshot::capture_and_save_streaming(vm, &virtio, &secondary_states, out)
            .map_err(|e| match e {
                snapshot::SnapshotStreamError::Hvf(h) => WorkerError::SnapshotCapture(h),
                snapshot::SnapshotStreamError::Io(io) => WorkerError::SnapshotSave {
                    path: out.to_string(),
                    source: io,
                },
            })?;
        // 0.7.6+ posix-fs sidecar: persist the (nodeid → host_path)
        // table next to the snapshot so warm restores don't lose it.
        // Best-effort — logs but does not fail the save.
        snapshot::capture_and_write_posix_fs_sidecar(out, posix_fs);
        // DAX sidecar: persist each session's slot table so the
        // restored guest can lazily re-bind on first stage-2 fault
        // instead of SIGBUSing. Same best-effort contract.
        snapshot::capture_and_write_dax_sidecar(out, dax_sessions);
        // smpark sidecar: pin smpark.ko's unpark-page GPA to the
        // snapshot so a RESTORED worker (which never boots, so never
        // sees the PL011 line that sets the process-global
        // SMPARK_STATE_GPA) can host-direct-unpark its secondaries with
        // the correct snapshot-scoped GPA. Best-effort; no-op when
        // single-vCPU / smpark absent.
        snapshot::capture_and_write_smpark_sidecar(
            snapshot.serial.smpark_state_gpa.load(Ordering::Acquire),
            out,
        );
        if !secondary_handles.is_empty() {
            coord.release_after_snapshot();
        }
        // Unpark — secondaries can resume normal execution. If we
        // parked successfully, this is essential: skipping it
        // leaves them stuck in WFI inside smpark_park_routine on
        // the resumed guest. If parking failed, this is a no-op
        // (the agent reports ok=false; we ignore the result).
        let total_us = t0.elapsed().as_micros();
        eprintln!("\n  snapshot ({reason}): total {total_us} us (capture+save streamed), RAM {} MiB (data {} MiB, zero {} MiB), mmio={} listeners={} -> {out}",
            write_stats.ram_bytes / (1024*1024),
            write_stats.ram_data_bytes / (1024*1024),
            write_stats.ram_zero_bytes / (1024*1024),
            virtio.devices.len(), virtio.vsock_listeners.len());
        Ok(true)
    };
    loop {
        if snapshot
            .stop_requested
            .is_some_and(|flag| flag.load(Ordering::SeqCst))
        {
            return Ok(DispatchExit::Stopped);
        }
        // Bake-then-pool readiness signal — fires the equivalent
        // of `take_snapshot` triggers (heartbeat / listener /
        // workload-parked / wall-clock) but returns a different
        // DispatchExit so the runner can hand control to the host
        // for warmup. Only one of `out_path` and
        // `bake_ready_signal` is set at any time.
        let bake_signal = snapshot.bake_ready_signal && snapshot.out_path.is_none();

        // Evaluate the priority-ordered snapshot triggers (heartbeat → pre-exec
        // → listener-ready → workload-parked → wall-clock) via the portable,
        // unit-tested predicate in `vmm::snapshot_trigger`. The ACTION it selects
        // — quiesce-to-WFI then capture, or signal bake-ready — is the
        // backend-specific part that stays here. (The per-trigger diagnostic
        // logs the inline chain printed collapse into take_snapshot's
        // "snapshot trigger (<reason>)" line.)
        let trigger_cfg = crate::vmm::snapshot_trigger::SnapshotTriggerConfig {
            at_heartbeat: snapshot.at_heartbeat,
            on_pre_exec: snapshot.on_pre_exec,
            on_listener: snapshot.on_listener,
            after_ms: snapshot.after_ms,
        };
        let signals = crate::vmm::snapshot_trigger::TriggerSignals {
            heartbeat_count: snapshot.serial.heartbeat_count.load(Ordering::SeqCst),
            pre_exec_ready: snapshot.serial.pre_exec_ready.load(Ordering::SeqCst),
            workload_parked: snapshot.serial.workload_parked.load(Ordering::SeqCst),
            listener_count: vsock.muxer().listener_count(),
            elapsed_ms: dispatch_start.elapsed().as_millis() as u64,
        };
        if let Some(fire) = trigger_cfg.evaluate(&signals) {
            // Quiesce-to-WFI for the triggers that can fire mid-activity
            // (heartbeat, listener-ready); the idle ones capture immediately.
            // `quiesced` guards against re-quiescing if a capture loops back.
            if fire.needs_quiesce && !quiesced && snapshot.quiesce_ms > 0 {
                eprintln!(
                    "  quiescing for {} ms before snapshot ({})...",
                    snapshot.quiesce_ms, fire.reason
                );
                if let Some(exit) =
                    quiesce_to_wfi(vcpu, bus, coord, dax_sessions, snapshot.quiesce_ms)?
                {
                    return Ok(exit);
                }
                quiesced = true;
            }
            if bake_signal {
                eprintln!("  bake-ready ({})", fire.reason);
                return Ok(DispatchExit::BakeReady);
            }
            if snapshot.out_path.is_some() && take_snapshot(fire.reason)? {
                return Ok(DispatchExit::SnapshotSaved);
            }
        }
        match vcpu_step(idx, vcpu, bus, coord, dax_sessions)? {
            StepExit::Canceled => return Ok(DispatchExit::Canceled),
            StepExit::Continue => {}
            StepExit::Stop => return Ok(DispatchExit::Stopped),
            StepExit::SystemOff => return Ok(DispatchExit::SystemOff),
        }
    }
}

#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
pub(crate) fn quiesce_to_wfi(
    vcpu: &Vcpu,
    bus: &MmioBus,
    coord: &VcpuCoordinator,
    dax_sessions: &[Arc<DaxSession>],
    wait_ms: u64,
) -> Result<Option<DispatchExit>, WorkerError> {
    vcpu.set_vtimer_mask(true)
        .map_err(WorkerError::VtimerMask)?;
    let h = vcpu.exit_token();
    let canceller = std::thread::Builder::new()
        .name("quiesce-canceller".into())
        // 0.7.43+: 64 KiB stack is plenty for a sleep + one
        // libc call. Default 2 MiB is wasteful for a thread
        // that lives <wait_ms milliseconds.
        .stack_size(64 * 1024)
        .spawn(move || {
            std::thread::sleep(std::time::Duration::from_millis(wait_ms));
            // Force this vCPU out of step() from the timer thread.
            VcpuHandle::force_exit(&[h]);
        })
        .map_err(WorkerError::QuiesceThreadSpawn)?;
    let result = loop {
        match vcpu_step(0, vcpu, bus, coord, dax_sessions)? {
            StepExit::Canceled => break Ok(None),
            StepExit::Continue => {}
            StepExit::Stop => break Ok(None),
            StepExit::SystemOff => break Ok(Some(DispatchExit::SystemOff)),
        }
    };
    let _ = canceller.join();
    let unmask = vcpu.set_vtimer_mask(false).map_err(WorkerError::VtimerMask);
    match (result, unmask) {
        (Err(e), _) => Err(e),
        (Ok(_), Err(e)) => Err(e),
        (Ok(exit), Ok(())) => Ok(exit),
    }
}

#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn dispatch_vcpu_inner(
    idx: u32,
    vcpu: &Vcpu,
    bus: &Arc<MmioBus>,
    coord: &Arc<VcpuCoordinator>,
    dax_sessions: &[Arc<DaxSession>],
) -> Result<(), WorkerError> {
    use std::sync::atomic::Ordering;
    loop {
        // Multi-vCPU snapshot rendezvous. When vcpu0 fires the
        // snapshot trigger, it sets `snapshot_request` + forces
        // each secondary out of `hv_vcpu_run` via
        // `hv_vcpus_exit`. The exit surfaces as
        // `StepExit::Canceled`. We must distinguish that from a
        // real shutdown: if `snapshot_request` is set, loop
        // around and let `maybe_pause_for_snapshot` capture
        // state; otherwise it's a true cancellation and we exit.
        coord
            .maybe_pause_for_snapshot(idx, vcpu)
            .map_err(WorkerError::SnapshotCapture)?;
        // Multi-vCPU cycle-restore rendezvous. Same shape as the
        // snapshot one above, but inverted: vcpu0 (the runner
        // pool-mode driver) sets `restore_request` after remapping
        // RAM + restoring its own state, and we apply our target
        // per-vCPU state on this OWNING thread (HVF requires it).
        // Without this, restore_on_release(true) + multi-vCPU
        // leaves secondaries on a stale trajectory and the next
        // softirq/IPI path panics.
        coord
            .maybe_apply_restore(idx, vcpu)
            .map_err(WorkerError::SnapshotCapture)?;
        match vcpu_step(idx, vcpu, bus, coord, dax_sessions)? {
            StepExit::Canceled => {
                if coord.snapshot_pause.is_paused() || coord.restore_request.load(Ordering::Acquire)
                {
                    // Forced exit by snapshot/restore trigger; loop
                    // back so the rendezvous point handles it.
                    continue;
                }
                return Ok(());
            }
            StepExit::Continue => {}
            StepExit::Stop | StepExit::SystemOff => return Ok(()),
        }
    }
}

/// Run one vcpu_run iteration and handle the exit.
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn vcpu_step(
    idx: u32,
    vcpu: &Vcpu,
    bus: &MmioBus,
    coord: &VcpuCoordinator,
    dax_sessions: &[Arc<DaxSession>],
) -> Result<StepExit, WorkerError> {
    use crate::vmm::coord::*;

    // 0.7.44+ removed exit_profile timing from the hot path. The
    // module's `enabled()` has been hard-coded `false` since 0.7.38
    // (SUPERMACHINE_EXIT_PROFILE deleted; `samply` / `cargo
    // flamegraph` supersede), so the `Instant::now()` + counter
    // bookkeeping at every vCPU exit was pure overhead — ~10k-50k
    // exits per guest boot × ~10-20 ns each. Module retained for
    // potential dev-time re-enable.
    let (esr, gpa) = match vcpu
        .step()
        .map_err(|source| WorkerError::VcpuRun { idx, source })?
    {
        VcpuExit::Exception {
            syndrome,
            phys_addr,
            ..
        } => (syndrome, phys_addr),
        VcpuExit::VTimerActivated => return Ok(StepExit::Continue),
        VcpuExit::Canceled => {
            eprintln!("  [vcpu-{idx}] canceled");
            return Ok(StepExit::Canceled);
        }
        VcpuExit::Unknown(v) => {
            eprintln!("  [vcpu-{idx}] unknown exit {v}");
            return Ok(StepExit::Stop);
        }
        // x86-only exits — HVF's aarch64 vCPU never produces these.
        VcpuExit::Io { .. } | VcpuExit::Mmio { .. } | VcpuExit::Halt => {
            eprintln!("  [vcpu-{idx}] unexpected x86 exit on HVF");
            return Ok(StepExit::Stop);
        }
    };
    let ec = (esr >> 26) & 0x3f;
    if ec == 0x24 {
        // DAX lazy-rehydrate: a cycle-restored VM loses the
        // host-side mmap + HVF stage-2 mapping for every DAX
        // slot that was active pre-snapshot — the snapshot
        // sidecar only persists the FUSE-level metadata. On
        // first guest access to a slot's GPA range we have
        // to re-`dax_map` + `hv_vm_map`, then let the vCPU
        // retry the load against the freshly-restored stage-
        // 2 entry. Without this hook the guest takes SIGBUS
        // on a tokio worker's first DAX touch (Prisma webhook
        // on cycle-restored Remix, the field repro).
        //
        // We DON'T advance PC on success — HVF re-executes
        // the trapped load and it now hits the present
        // mapping. This is cheaper and more correct than
        // emulating the instruction (we'd have to decode the
        // full aarch64 load/store space; the architecture
        // does it for free on retry).
        let mut dax_consumed = false;
        for sess in dax_sessions {
            match sess.handle_stage2_fault(gpa) {
                Ok(true) => {
                    dax_consumed = true;
                    break;
                }
                Ok(false) => continue,
                Err(e) => {
                    eprintln!(
                        "[vcpu-{idx}] dax stage-2 fault gpa={gpa:#x} re-map FAILED \
                                 errno={e}; falling through to MMIO dispatch (will return 0)"
                    );
                    break;
                }
            }
        }
        if !dax_consumed {
            handle_data_abort_vcpu(vcpu, bus, esr, gpa);
        }
        // dax_consumed: vCPU resumes at same PC; HVF retries
        // the load against the now-mapped GPA.
    } else if ec == 0x20 || ec == 0x21 {
        // Instruction-abort (lower EL = 0x20, same EL = 0x21).
        // Most common case in the field: Prisma's tokio worker
        // jumps into a code page that was DAX-mapped pre-
        // snapshot but never paged-in. Without lazy rehydrate
        // the guest takes a fatal abort here. Same return-
        // without-advance contract as the data-abort branch.
        let mut dax_consumed = false;
        for sess in dax_sessions {
            match sess.handle_stage2_fault(gpa) {
                Ok(true) => {
                    dax_consumed = true;
                    break;
                }
                Ok(false) => continue,
                Err(e) => {
                    eprintln!(
                        "[vcpu-{idx}] dax instruction-abort gpa={gpa:#x} re-map FAILED \
                                 errno={e}; falling through to unhandled-EC path"
                    );
                    break;
                }
            }
        }
        if !dax_consumed {
            let pc = vcpu.get_core(CoreReg::Pc).unwrap_or(0);
            eprintln!(
                "  [vcpu-{idx}] unhandled instruction-abort EC={ec:#x} ESR=0x{esr:x} \
                         PC=0x{pc:x} FAR=0x{gpa:x}"
            );
            return Ok(StepExit::Stop);
        }
    } else if ec == 0x16 {
        // HVC: PSCI dispatch.
        let fid = vcpu.get_core(CoreReg::X(0)).unwrap_or(0) as u32;
        let ret: i64 = match fid {
            PSCI_VERSION => 0x10000, // PSCI 1.0
            PSCI_FEATURES => {
                let q = vcpu.get_core(CoreReg::X(1)).unwrap_or(0) as u32;
                match q {
                    PSCI_VERSION | PSCI_CPU_ON | PSCI_CPU_OFF | PSCI_AFFINITY_INFO
                    | PSCI_FEATURES | PSCI_SYSTEM_OFF | PSCI_SYSTEM_RESET => 0,
                    _ => PSCI_NOT_SUPPORTED,
                }
            }
            PSCI_CPU_ON => {
                let target = vcpu.get_core(CoreReg::X(1)).unwrap_or(0) as u32;
                let entry = vcpu.get_core(CoreReg::X(2)).unwrap_or(0);
                let ctx_id = vcpu.get_core(CoreReg::X(3)).unwrap_or(0);
                eprintln!("  [vcpu-{idx}] PSCI CPU_ON target={target} entry=0x{entry:x}");
                coord.cpu_on(target, entry, ctx_id)
            }
            PSCI_CPU_OFF => {
                // The PSCI spec says CPU_OFF doesn't return; stop dispatching.
                eprintln!("  [vcpu-{idx}] PSCI CPU_OFF");
                return Ok(StepExit::Stop);
            }
            PSCI_AFFINITY_INFO => {
                let target = vcpu.get_core(CoreReg::X(1)).unwrap_or(0) as u32;
                coord.affinity_info(target)
            }
            PSCI_SYSTEM_OFF | PSCI_SYSTEM_RESET => {
                eprintln!("  [vcpu-{idx}] PSCI SYSTEM_OFF/RESET - exiting");
                return Ok(StepExit::SystemOff);
            }
            _ => PSCI_NOT_SUPPORTED,
        };
        vcpu.set_core(CoreReg::X(0), ret as u64).ok();
    } else if ec == 0x18 {
        let pc = vcpu.get_core(CoreReg::Pc).unwrap_or(0);
        let iss = esr & 0x01ff_ffff;
        let rt = ((iss >> 5) & 0x1f) as u32;
        let is_read = (iss & 1) != 0;
        if is_read && rt < 31 {
            vcpu.set_core(CoreReg::X(rt as u8), 0).ok();
        }
        vcpu.set_core(CoreReg::Pc, pc + 4).ok();
    } else if ec == 0x17 {
        let pc = vcpu.get_core(CoreReg::Pc).unwrap_or(0);
        vcpu.set_core(CoreReg::X(0), (-1i64) as u64).ok();
        vcpu.set_core(CoreReg::Pc, pc + 4).ok();
    } else {
        let pc = vcpu.get_core(CoreReg::Pc).unwrap_or(0);
        eprintln!("  [vcpu-{idx}] unhandled EC={ec:#x} ESR=0x{esr:x} PC=0x{pc:x}");
        return Ok(StepExit::Stop);
    }
    Ok(StepExit::Continue)
}

/// ESR_EL2 ISS for "Data Abort, lower EL" (EC=0x24, IFSC=0x10).
/// Operates on any vCPU so secondary vCPU threads can dispatch their own MMIO
/// traps.
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn handle_data_abort_vcpu(vcpu: &Vcpu, bus: &MmioBus, esr: u64, gpa: u64) {
    let pc = vcpu.get_core(CoreReg::Pc).unwrap_or(0);
    let far = gpa;
    let iss = esr & 0x01ff_ffff;
    let isv = ((iss >> 24) & 1) != 0;
    if !isv {
        vcpu.set_core(CoreReg::Pc, pc + 4).ok();
        return;
    }
    let sas = ((iss >> 22) & 0x3) as u8;
    let size: u8 = 1 << sas;
    let srt = ((iss >> 16) & 0x1f) as u32;
    let wnr = ((iss >> 6) & 1) != 0;

    // 0.7.44+ exit_profile::record removed from hot path — see
    // `vcpu_step` comment. Used to wrap bus.{read,write} in
    // Instant::now() + record() per MMIO exit.
    if wnr {
        let val = vcpu.get_x(srt).unwrap_or(0);
        let handled = bus.write(far, val, size);
        if !handled && std::env::var("SUPERMACHINE_TRACE").is_ok() {
            eprintln!("MMIO W {far:#x} = {val:#x} sz={size} (unhandled)");
        }
    } else {
        let val = bus.read(far, size);
        if val.is_none() && std::env::var("SUPERMACHINE_TRACE").is_ok() {
            eprintln!("MMIO R {far:#x} sz={size} (unhandled, returning 0)");
        }
        let val = val.unwrap_or(0);
        if srt < 31 {
            vcpu.set_x(srt, val).ok();
        }
    }
    vcpu.set_core(CoreReg::Pc, pc + 4).ok();
}
