//! Pool-worker supervisor protocol support.
//!
//! The CLI performs the initial connect/READY/RESTORE handshake. Once a VM is
//! running, this module owns the persistent supervisor reader thread and the
//! small shared state used to request warm restores.
//!
//! 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::io::Write;
use std::os::unix::net::UnixStream;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc;
use std::sync::{Arc, Mutex};
use std::thread::JoinHandle;
use std::time::Duration;

use crate::vmm::resources::VmResources;
use crate::vmm::runner::{run, RunError, RunOptions, RunReport};

type TransportIdleCheck = Arc<dyn Fn() -> bool + Send + Sync + 'static>;

#[derive(Debug)]
pub struct RestoreRequest {
    pub path: String,
    pub egress_policy: Option<String>,
}

/// Mid-flight snapshot capture request — pause the running VM,
/// capture full state to `out_path`, signal completion. Used by
/// the public [`crate::Vm::snapshot`] API to mint a new
/// snapshot from a warmed-up VM (the "rustc-warm snapshot"
/// pattern: boot, populate target/, snapshot, reuse).
#[derive(Debug)]
pub struct SnapshotRequest {
    pub out_path: String,
    /// `Sync` (default): capture+save under guest pause, return
    /// `DONE_SNAPSHOT` only after the file is on disk.
    /// `Async`: capture into a compact in-memory buffer (~50 ms
    /// for ~100 MiB of non-zero pages), spawn a background
    /// thread that writes the buffer to disk via `.partial`+
    /// rename, return `DONE_SNAPSHOT_ASYNC` immediately. The
    /// runner tracks in-flight async saves; `QUIT` waits for
    /// them to drain. Used by the pipelined-bake flow so the
    /// base snapshot save overlaps with the warmup workload.
    pub mode: SnapshotMode,
    /// Optional base-snapshot path hint for the differential
    /// save path. When set AND mode == Sync AND the runner has
    /// the matching base in memory (from a prior `SNAPSHOT_ASYNC
    /// <base>`), the runner uses APFS clonefile + diff-pwrite
    /// instead of the streaming sparse-write path, cutting warm
    /// save from ~300 ms to ~30 ms on typical workloads. The
    /// runner waits for the in-flight base save to finish
    /// before opening the clone.
    ///
    /// Falls through to plain streaming sync save if:
    ///   * mode != Sync
    ///   * base_path is None
    ///   * the runner has no in-flight save matching base_path
    ///   * `clonefile(2)` returns EXDEV (different filesystem)
    ///   * warm meta overflows base's `ram_offset` slack (rare)
    pub base_path: Option<String>,
    /// `inline` token. When set (Sync + base diff path), the runner waits for
    /// the diff save to land before returning `DONE_SNAPSHOT`, so the caller
    /// can read the file immediately — without giving up the in-memory base
    /// retention that lets a *chain* of diffs (the in-VM builder's per-layer
    /// snapshots) each diff against the previous in memory. Default `false`
    /// keeps the existing fire-and-forget bg-save behavior.
    pub inline_diff: bool,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SnapshotMode {
    Sync,
    Async,
}

#[derive(Clone, Debug)]
pub struct SnapshotResult {
    pub bytes_written: u64,
    pub capture_us: u128,
    pub save_us: u128,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct PoolRestoreResult {
    pub restore_us: u128,
    pub host_port: Option<u16>,
    pub timings: WarmRestoreTimings,
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct WarmRestoreTimings {
    pub reset_vsock_us: u128,
    pub remap_cow_us: u128,
    pub load_meta_us: u128,
    pub restore_snapshot_us: u128,
    pub ram_copy_us: u128,
    pub gic_restore_us: u128,
    pub vcpu_restore_us: u128,
    pub vtimer_offset_us: u128,
    pub mmio_restore_us: u128,
    pub listener_restore_us: u128,
}

#[derive(Clone)]
pub struct PoolHandle {
    cmd_tx: mpsc::Sender<PoolCommand>,
    request_lock: Arc<Mutex<()>>,
    busy: Arc<AtomicBool>,
}

pub struct PoolWorker {
    cmd_rx: mpsc::Receiver<PoolCommand>,
    active_done: Arc<Mutex<Option<mpsc::Sender<PoolRestoreResult>>>>,
    busy: Arc<AtomicBool>,
}

pub struct WarmPool {
    handle: PoolHandle,
    vm_thread: Option<JoinHandle<Result<RunReport, RunError>>>,
}

pub struct PoolControl {
    pause_pending: Arc<AtomicBool>,
    restore_req: Arc<Mutex<Option<RestoreRequest>>>,
    snapshot_req: Arc<Mutex<Option<SnapshotRequest>>>,
    snapshot_done: Arc<Mutex<Option<mpsc::Sender<Result<SnapshotResult, String>>>>>,
    quit: Arc<AtomicBool>,
    writer: Option<UnixStream>,
    /// Same Mutex-protected supervisor-socket writer used by the
    /// SNAPSHOT response path; cloned out to PoolControl so the
    /// runner can also write `BAKE_READY` (bake-then-pool flow)
    /// without racing the supervisor's `DONE_SNAPSHOT` writes.
    supervisor_writer: Option<Arc<Mutex<Option<UnixStream>>>>,
    active_done: Option<Arc<Mutex<Option<mpsc::Sender<PoolRestoreResult>>>>>,
    busy: Option<Arc<AtomicBool>>,
}

#[derive(Debug)]
pub enum PoolError {
    InitialDoneClone(std::io::Error),
    InitialDoneWrite(std::io::Error),
    ReaderClone(std::io::Error),
    ThreadSpawn(std::io::Error),
}

#[derive(Debug, PartialEq, Eq)]
pub enum PoolClientError {
    CommandClosed,
    CompletionClosed,
    CompletionTimeout,
    RestoreInFlight,
    /// `PoolHandle::snapshot` succeeded in the round-trip but the
    /// runner reported a capture/save error. Carries the runner's
    /// error message.
    SnapshotFailed(String),
}

#[derive(Debug)]
pub enum WarmPoolError {
    AlreadyJoined,
    Pool(PoolClientError),
    PoolWorkerAlreadySet,
    Run(RunError),
    ThreadJoin,
    ThreadSpawn(std::io::Error),
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn warm_pool_restore_fails_when_worker_is_dropped() {
        let (handle, worker) = warm_pool();
        drop(worker);
        assert_eq!(
            handle.restore("snap.sm"),
            Err(PoolClientError::CommandClosed)
        );
    }

    #[test]
    fn warm_pool_quit_fails_when_worker_is_dropped() {
        let (handle, worker) = warm_pool();
        drop(worker);
        assert_eq!(handle.quit(), Err(PoolClientError::CommandClosed));
    }

    #[test]
    fn warm_pool_restore_receives_per_request_completion() {
        let (handle, worker) = warm_pool();
        let restore = std::thread::spawn(move || handle.restore("snap.sm"));

        match worker.cmd_rx.recv().unwrap() {
            PoolCommand::Restore { request, done_tx } => {
                assert_eq!(request.path, "snap.sm");
                done_tx
                    .send(PoolRestoreResult {
                        restore_us: 123,
                        host_port: Some(61820),
                        timings: WarmRestoreTimings {
                            reset_vsock_us: 1,
                            remap_cow_us: 2,
                            load_meta_us: 3,
                            restore_snapshot_us: 4,
                            ram_copy_us: 5,
                            gic_restore_us: 6,
                            vcpu_restore_us: 7,
                            vtimer_offset_us: 8,
                            mmio_restore_us: 9,
                            listener_restore_us: 10,
                        },
                    })
                    .unwrap();
            }
            PoolCommand::Quit => panic!("unexpected quit command"),
            PoolCommand::Snapshot { .. } => panic!("unexpected snapshot command"),
        }

        assert_eq!(
            restore.join().unwrap(),
            Ok(PoolRestoreResult {
                restore_us: 123,
                host_port: Some(61820),
                timings: WarmRestoreTimings {
                    reset_vsock_us: 1,
                    remap_cow_us: 2,
                    load_meta_us: 3,
                    restore_snapshot_us: 4,
                    ram_copy_us: 5,
                    gic_restore_us: 6,
                    vcpu_restore_us: 7,
                    vtimer_offset_us: 8,
                    mmio_restore_us: 9,
                    listener_restore_us: 10,
                },
            })
        );
    }

    #[test]
    fn warm_pool_restore_timeout_keeps_request_in_flight() {
        let (handle, _worker) = warm_pool();
        assert_eq!(
            handle.restore_timeout("snap.sm", Duration::ZERO),
            Err(PoolClientError::CompletionTimeout)
        );
        assert_eq!(
            handle.restore_timeout("snap.sm", Duration::ZERO),
            Err(PoolClientError::RestoreInFlight)
        );
    }

    #[test]
    fn warm_pool_manager_rejects_existing_worker() {
        let (_handle, worker) = warm_pool();
        let result = WarmPool::start(
            VmResources::from_snapshot("snap.sm"),
            RunOptions {
                pool_worker: Some(worker),
                ..RunOptions::default()
            },
        );

        assert!(matches!(result, Err(WarmPoolError::PoolWorkerAlreadySet)));
    }

    // ── busy-flag state machine: restore ⇄ snapshot serialization ──────

    #[test]
    fn restore_command_closed_resets_busy() {
        // A failed send must NOT leave the slot wedged busy — otherwise a
        // transient channel error would permanently reject every retry as
        // RestoreInFlight.
        let (handle, worker) = warm_pool();
        drop(worker);
        assert_eq!(
            handle.restore("snap.sm"),
            Err(PoolClientError::CommandClosed)
        );
        assert!(
            !handle.busy.load(Ordering::SeqCst),
            "busy must be cleared after a CommandClosed send failure"
        );
    }

    #[test]
    fn snapshot_command_closed_resets_busy() {
        let (handle, worker) = warm_pool();
        drop(worker);
        assert!(
            matches!(
                handle.snapshot("out.sm"),
                Err(PoolClientError::CommandClosed)
            ),
            "dropped worker → CommandClosed"
        );
        assert!(
            !handle.busy.load(Ordering::SeqCst),
            "busy cleared on closed"
        );
    }

    #[test]
    fn snapshot_success_round_trips_and_clears_busy() {
        let (handle, worker) = warm_pool();
        let h = std::thread::spawn(move || handle.snapshot("out.sm"));

        match worker.cmd_rx.recv().unwrap() {
            PoolCommand::Snapshot { request, done_tx } => {
                assert_eq!(request.out_path, "out.sm");
                assert!(matches!(request.mode, SnapshotMode::Sync));
                done_tx
                    .send(Ok(SnapshotResult {
                        bytes_written: 4096,
                        capture_us: 11,
                        save_us: 22,
                    }))
                    .unwrap();
            }
            other => panic!(
                "expected Snapshot, got {:?}",
                std::mem::discriminant(&other)
            ),
        }

        let r = h.join().unwrap().expect("snapshot ok");
        assert_eq!(r.bytes_written, 4096);
        assert_eq!(r.save_us, 22);
    }

    #[test]
    fn snapshot_runner_error_maps_to_snapshot_failed() {
        let (handle, worker) = warm_pool();
        let h = std::thread::spawn(move || handle.snapshot("out.sm"));

        match worker.cmd_rx.recv().unwrap() {
            PoolCommand::Snapshot { done_tx, .. } => {
                done_tx.send(Err("disk full".to_string())).unwrap();
            }
            _ => panic!("expected Snapshot"),
        }
        assert!(
            matches!(h.join().unwrap(), Err(PoolClientError::SnapshotFailed(ref m)) if m == "disk full"),
            "runner error maps to SnapshotFailed(msg)"
        );
    }

    #[test]
    fn snapshot_timeout_clears_busy_so_a_retry_is_not_in_flight() {
        // Snapshot timeout (unlike restore timeout) frees the slot: the
        // runner-side capture is abandoned, so a fresh snapshot may retry.
        // Contrast warm_pool_restore_timeout_keeps_request_in_flight.
        let (handle, _worker) = warm_pool();
        assert!(
            matches!(
                handle.snapshot_timeout("out.sm", Duration::ZERO),
                Err(PoolClientError::CompletionTimeout)
            ),
            "first snapshot times out"
        );
        assert!(
            !handle.busy.load(Ordering::SeqCst),
            "snapshot timeout clears busy"
        );
        assert!(
            matches!(
                handle.snapshot_timeout("out.sm", Duration::ZERO),
                Err(PoolClientError::CompletionTimeout)
            ),
            "a second snapshot retries (NOT RestoreInFlight)"
        );
    }

    #[test]
    fn busy_flag_serializes_restore_against_snapshot() {
        // restore timeout wedges the slot busy (terminal). A snapshot then
        // shares the SAME busy flag and must be rejected as in-flight —
        // the runner can only do one pause-and-act at a time.
        let (handle, _worker) = warm_pool();
        assert_eq!(
            handle.restore_timeout("snap.sm", Duration::ZERO),
            Err(PoolClientError::CompletionTimeout)
        );
        assert!(
            handle.busy.load(Ordering::SeqCst),
            "restore timeout keeps busy"
        );
        assert!(
            matches!(
                handle.snapshot("out.sm"),
                Err(PoolClientError::RestoreInFlight)
            ),
            "snapshot must observe the restore's busy slot"
        );
    }
}

enum PoolCommand {
    Restore {
        request: RestoreRequest,
        done_tx: mpsc::Sender<PoolRestoreResult>,
    },
    Snapshot {
        request: SnapshotRequest,
        done_tx: mpsc::Sender<Result<SnapshotResult, String>>,
    },
    Quit,
}

pub fn warm_pool() -> (PoolHandle, PoolWorker) {
    let (cmd_tx, cmd_rx) = mpsc::channel();
    let active_done = Arc::new(Mutex::new(None));
    let busy = Arc::new(AtomicBool::new(false));
    (
        PoolHandle {
            cmd_tx,
            request_lock: Arc::new(Mutex::new(())),
            busy: busy.clone(),
        },
        PoolWorker {
            cmd_rx,
            active_done,
            busy,
        },
    )
}

impl PoolHandle {
    pub fn restore(&self, path: impl Into<String>) -> Result<PoolRestoreResult, PoolClientError> {
        self.restore_request(RestoreRequest {
            path: path.into(),
            egress_policy: None,
        })
    }

    pub fn restore_with_egress_policy(
        &self,
        path: impl Into<String>,
        egress_policy: impl Into<String>,
    ) -> Result<PoolRestoreResult, PoolClientError> {
        self.restore_request(RestoreRequest {
            path: path.into(),
            egress_policy: Some(egress_policy.into()),
        })
    }

    pub fn restore_timeout(
        &self,
        path: impl Into<String>,
        timeout: Duration,
    ) -> Result<PoolRestoreResult, PoolClientError> {
        self.restore_request_timeout(
            RestoreRequest {
                path: path.into(),
                egress_policy: None,
            },
            timeout,
        )
    }

    pub fn restore_with_egress_policy_timeout(
        &self,
        path: impl Into<String>,
        egress_policy: impl Into<String>,
        timeout: Duration,
    ) -> Result<PoolRestoreResult, PoolClientError> {
        self.restore_request_timeout(
            RestoreRequest {
                path: path.into(),
                egress_policy: Some(egress_policy.into()),
            },
            timeout,
        )
    }

    pub fn quit(&self) -> Result<(), PoolClientError> {
        self.cmd_tx
            .send(PoolCommand::Quit)
            .map_err(|_| PoolClientError::CommandClosed)
    }

    /// Capture a snapshot of the currently running VM to
    /// `out_path`. Pauses the VM via the same mechanism as
    /// [`PoolHandle::restore`], runs `capture_snapshot` +
    /// `save_to_file` on the runner thread, signals completion.
    ///
    /// Used by the public [`crate::Vm::snapshot`] API.
    pub fn snapshot(&self, out_path: impl Into<String>) -> Result<SnapshotResult, PoolClientError> {
        self.snapshot_request_inner(
            SnapshotRequest {
                out_path: out_path.into(),
                mode: SnapshotMode::Sync,
                base_path: None,
                inline_diff: false,
            },
            None,
        )
    }

    pub fn snapshot_timeout(
        &self,
        out_path: impl Into<String>,
        timeout: Duration,
    ) -> Result<SnapshotResult, PoolClientError> {
        self.snapshot_request_inner(
            SnapshotRequest {
                out_path: out_path.into(),
                mode: SnapshotMode::Sync,
                base_path: None,
                inline_diff: false,
            },
            Some(timeout),
        )
    }

    fn snapshot_request_inner(
        &self,
        request: SnapshotRequest,
        timeout: Option<Duration>,
    ) -> Result<SnapshotResult, PoolClientError> {
        // Serialize on the same lock restore uses — the runner
        // can only handle one pause-and-do-thing at a time.
        let _lock = self
            .request_lock
            .lock()
            .map_err(|_| PoolClientError::CommandClosed)?;
        if self.busy.swap(true, Ordering::SeqCst) {
            return Err(PoolClientError::RestoreInFlight);
        }
        let (done_tx, done_rx) = mpsc::channel();
        self.cmd_tx
            .send(PoolCommand::Snapshot { request, done_tx })
            .map_err(|_| {
                self.busy.store(false, Ordering::SeqCst);
                PoolClientError::CommandClosed
            })?;
        let result = match timeout {
            Some(t) => done_rx.recv_timeout(t).map_err(|e| match e {
                mpsc::RecvTimeoutError::Timeout => PoolClientError::CompletionTimeout,
                mpsc::RecvTimeoutError::Disconnected => PoolClientError::CompletionClosed,
            }),
            None => done_rx
                .recv()
                .map_err(|_| PoolClientError::CompletionClosed),
        };
        self.busy.store(false, Ordering::SeqCst);
        match result? {
            Ok(r) => Ok(r),
            Err(msg) => Err(PoolClientError::SnapshotFailed(msg)),
        }
    }

    fn restore_request(
        &self,
        request: RestoreRequest,
    ) -> Result<PoolRestoreResult, PoolClientError> {
        self.restore_request_inner(request, None)
    }

    fn restore_request_timeout(
        &self,
        request: RestoreRequest,
        timeout: Duration,
    ) -> Result<PoolRestoreResult, PoolClientError> {
        self.restore_request_inner(request, Some(timeout))
    }

    fn restore_request_inner(
        &self,
        request: RestoreRequest,
        timeout: Option<Duration>,
    ) -> Result<PoolRestoreResult, PoolClientError> {
        let _guard = self.request_lock.lock().unwrap();
        if self.busy.swap(true, Ordering::SeqCst) {
            return Err(PoolClientError::RestoreInFlight);
        }

        let (done_tx, done_rx) = mpsc::channel();
        self.cmd_tx
            .send(PoolCommand::Restore { request, done_tx })
            .map_err(|_| {
                self.busy.store(false, Ordering::SeqCst);
                PoolClientError::CommandClosed
            })?;

        let result = match timeout {
            Some(timeout) => match done_rx.recv_timeout(timeout) {
                Ok(result) => Ok(result),
                Err(mpsc::RecvTimeoutError::Timeout) => Err(PoolClientError::CompletionTimeout),
                Err(mpsc::RecvTimeoutError::Disconnected) => Err(PoolClientError::CompletionClosed),
            },
            None => done_rx
                .recv()
                .map_err(|_| PoolClientError::CompletionClosed),
        };

        if !matches!(result, Err(PoolClientError::CompletionTimeout)) {
            self.busy.store(false, Ordering::SeqCst);
        }
        result
    }
}

impl WarmPool {
    pub fn start(resources: VmResources, mut options: RunOptions) -> Result<Self, WarmPoolError> {
        if options.pool_worker.is_some() {
            return Err(WarmPoolError::PoolWorkerAlreadySet);
        }

        let (handle, worker) = warm_pool();
        options.pool_worker = Some(worker);
        let vm_thread = std::thread::Builder::new()
            .name("warm-pool-vm".into())
            .spawn(move || {
                set_vm_thread_qos();
                run(&resources, options)
            })
            .map_err(WarmPoolError::ThreadSpawn)?;

        Ok(Self {
            handle,
            vm_thread: Some(vm_thread),
        })
    }

    pub fn handle(&self) -> PoolHandle {
        self.handle.clone()
    }

    pub fn restore(&self, path: impl Into<String>) -> Result<PoolRestoreResult, PoolClientError> {
        self.handle.restore(path)
    }

    pub fn restore_timeout(
        &self,
        path: impl Into<String>,
        timeout: Duration,
    ) -> Result<PoolRestoreResult, PoolClientError> {
        self.handle.restore_timeout(path, timeout)
    }

    pub fn restore_with_egress_policy(
        &self,
        path: impl Into<String>,
        egress_policy: impl Into<String>,
    ) -> Result<PoolRestoreResult, PoolClientError> {
        self.handle.restore_with_egress_policy(path, egress_policy)
    }

    pub fn restore_with_egress_policy_timeout(
        &self,
        path: impl Into<String>,
        egress_policy: impl Into<String>,
        timeout: Duration,
    ) -> Result<PoolRestoreResult, PoolClientError> {
        self.handle
            .restore_with_egress_policy_timeout(path, egress_policy, timeout)
    }

    /// Capture a snapshot of the running VM to `out_path`.
    /// Used by [`crate::Vm::snapshot`].
    pub fn snapshot(&self, out_path: impl Into<String>) -> Result<SnapshotResult, PoolClientError> {
        self.handle.snapshot(out_path)
    }

    pub fn snapshot_timeout(
        &self,
        out_path: impl Into<String>,
        timeout: Duration,
    ) -> Result<SnapshotResult, PoolClientError> {
        self.handle.snapshot_timeout(out_path, timeout)
    }

    pub fn shutdown(mut self) -> Result<RunReport, WarmPoolError> {
        let quit = self.handle.quit();
        let joined = self.join();
        match joined {
            Ok(report) => {
                quit.map_err(WarmPoolError::Pool)?;
                Ok(report)
            }
            Err(e) => Err(e),
        }
    }

    pub fn join(&mut self) -> Result<RunReport, WarmPoolError> {
        let vm_thread = self.vm_thread.take().ok_or(WarmPoolError::AlreadyJoined)?;
        vm_thread
            .join()
            .map_err(|_| WarmPoolError::ThreadJoin)?
            .map_err(WarmPoolError::Run)
    }
}

#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn set_vm_thread_qos() {
    const QOS_CLASS_USER_INTERACTIVE: u32 = 0x21;
    unsafe extern "C" {
        fn pthread_set_qos_class_self_np(qos_class: u32, relative_priority: i32) -> i32;
    }

    // Best-effort: QoS is an optimization for the managed library worker thread,
    // not a correctness dependency.
    unsafe {
        let _ = pthread_set_qos_class_self_np(QOS_CLASS_USER_INTERACTIVE, 0);
    }
}

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

impl fmt::Display for PoolError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            PoolError::InitialDoneClone(e) => {
                write!(f, "clone pool socket for initial DONE: {e}")
            }
            PoolError::InitialDoneWrite(e) => write!(f, "write initial pool DONE: {e}"),
            PoolError::ReaderClone(e) => write!(f, "clone pool socket for reader: {e}"),
            PoolError::ThreadSpawn(e) => write!(f, "spawn pool supervisor reader: {e}"),
        }
    }
}

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

impl fmt::Display for PoolClientError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            PoolClientError::CommandClosed => write!(f, "warm pool command channel closed"),
            PoolClientError::CompletionClosed => write!(f, "warm pool completion channel closed"),
            PoolClientError::CompletionTimeout => write!(f, "warm pool completion timed out"),
            PoolClientError::RestoreInFlight => write!(f, "warm pool restore already in flight"),
            PoolClientError::SnapshotFailed(msg) => write!(f, "snapshot failed: {msg}"),
        }
    }
}

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

impl fmt::Display for WarmPoolError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            WarmPoolError::AlreadyJoined => write!(f, "warm pool VM thread already joined"),
            WarmPoolError::Pool(e) => write!(f, "{e}"),
            WarmPoolError::PoolWorkerAlreadySet => {
                write!(f, "RunOptions already contains a pool worker")
            }
            WarmPoolError::Run(e) => write!(f, "{e}"),
            WarmPoolError::ThreadJoin => write!(f, "warm pool VM thread panicked"),
            WarmPoolError::ThreadSpawn(e) => write!(f, "spawn warm pool VM thread: {e}"),
        }
    }
}

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

impl PoolControl {
    #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
    pub fn start(
        sock: Option<&UnixStream>,
        worker: Option<PoolWorker>,
        initial_restore_us: Option<u128>,
        initial_host_port: Option<u16>,
        vcpu0_handle: crate::hypervisor::ActiveVcpuHandle,
        transport_idle: Option<TransportIdleCheck>,
        // This VM's per-VM PL011 state (`Vmm::serial`). The supervisor
        // reader consults its `smpark_state_gpa` (bake host-direct
        // unpark fallback) and pushes RX bytes through it (PUSH_RX);
        // the kcache marker watcher polls its `kcache_ready`.
        serial: Arc<crate::devices::serial::SerialState>,
    ) -> Result<Self, PoolError> {
        let pause_pending = Arc::new(AtomicBool::new(false));
        let restore_req = Arc::new(Mutex::new(None));
        let snapshot_req = Arc::new(Mutex::new(None));
        let snapshot_done = Arc::new(Mutex::new(None));
        let quit = Arc::new(AtomicBool::new(false));
        let mut writer = None;
        let mut active_done = None;
        let mut busy = None;

        // Shared writer end used by BOTH the initial-DONE write
        // (above) and the supervisor reader (for SNAPSHOT
        // responses). Wrapped in a Mutex so the two writers
        // don't race on the byte stream.
        let snapshot_done_writer: Arc<Mutex<Option<UnixStream>>> = Arc::new(Mutex::new(None));
        if let Some(sock) = sock {
            // Stash a write half for use by the SNAPSHOT response
            // path even when initial_restore_us is None.
            let writer_clone = sock.try_clone().map_err(PoolError::InitialDoneClone)?;
            *snapshot_done_writer.lock().unwrap() = Some(writer_clone);
            // ALWAYS clone a writer for cycle-restore DONE responses.
            // This was previously gated on `initial_restore_us.is_some()`
            // because the same stream was reused for both the initial
            // DONE (after spawn-time restore) and subsequent cycle-
            // restore DONEs. That gating broke the bake-then-pool /
            // warm-handoff path: the worker has no initial restore, so
            // `writer` stayed None, and `complete_restore` silently
            // skipped writing DONE — the library's `send_restore` then
            // hung forever waiting on `read_line`. The fix splits the
            // two: we always have a writer for cycle-restore DONEs,
            // and only the initial DONE write is conditional.
            let cycle_writer = sock.try_clone().map_err(PoolError::InitialDoneClone)?;
            writer = Some(cycle_writer);
            if let Some(us) = initial_restore_us {
                let mut initial_s = sock.try_clone().map_err(PoolError::InitialDoneClone)?;
                write_done(
                    &mut initial_s,
                    us,
                    initial_host_port,
                    WarmRestoreTimings::default(),
                )
                .map_err(PoolError::InitialDoneWrite)?;
                // Note: `writer` is the cycle-restore writer (above);
                // `initial_s` is consumed by the initial-DONE write
                // and dropped here. They're separate clones of the
                // same UnixStream, so writes don't race (the kernel
                // serializes per-fd writes ≤ PIPE_BUF, and our DONE
                // lines are well under that).
                let _ = initial_s;
            }

            let pp = pause_pending.clone();
            let rr = restore_req.clone();
            let sr = snapshot_req.clone();
            let sd = snapshot_done.clone();
            let sdw = snapshot_done_writer.clone();
            let qf = quit.clone();
            let transport_idle = transport_idle.clone();
            let read_half = sock.try_clone().map_err(PoolError::ReaderClone)?;
            let serial_reader = serial.clone();
            std::thread::Builder::new()
                .name("pool-supervisor-reader".into())
                .spawn(move || {
                    read_supervisor(
                        read_half,
                        pp,
                        rr,
                        sr,
                        sd,
                        sdw,
                        qf,
                        vcpu0_handle,
                        transport_idle,
                        serial_reader,
                    )
                })
                .map_err(PoolError::ThreadSpawn)?;

            // 0.7.43+ kernel-boot cache marker watcher. Polls the
            // `KCACHE_READY` atomic set by the PL011 line scanner
            // when init-oci prints the "[SUPERMACHINE-KCACHE]
            // ready" marker, and emits `MARKER_KCACHE_READY` on
            // the supervisor socket so the bake driver can react
            // (take a cache snapshot, then push the resume token).
            //
            // Single-shot per worker lifetime: emits at most once.
            // The bake driver re-spawns the worker for each bake
            // (no in-process state carryover).
            //
            // 100 ms poll interval — the marker fires once per
            // bake within ~100-500 ms of boot, so the per-poll
            // latency is in the noise relative to the bake's
            // total wall-clock.
            let watcher_writer = snapshot_done_writer.clone();
            let watcher_quit = quit.clone();
            let watcher_serial = serial.clone();
            std::thread::Builder::new()
                .name("pool-kcache-marker-watcher".into())
                .stack_size(64 * 1024)
                .spawn(move || {
                    use std::io::Write;
                    use std::sync::atomic::Ordering as O;
                    loop {
                        if watcher_quit.load(O::SeqCst) {
                            break;
                        }
                        if watcher_serial.kcache_ready.load(O::SeqCst) {
                            if let Ok(mut w) = watcher_writer.lock() {
                                if let Some(s) = w.as_mut() {
                                    let _ = writeln!(s, "MARKER_KCACHE_READY");
                                    let _ = s.flush();
                                }
                            }
                            // Reset so a subsequent fire (shouldn't
                            // happen, but defensive) re-emits.
                            watcher_serial.kcache_ready.store(false, O::SeqCst);
                            break;
                        }
                        std::thread::sleep(std::time::Duration::from_millis(100));
                    }
                })
                .map_err(PoolError::ThreadSpawn)?;
        }

        if let Some(worker) = worker {
            let pp = pause_pending.clone();
            let rr = restore_req.clone();
            let sr = snapshot_req.clone();
            let sd = snapshot_done.clone();
            let qf = quit.clone();
            let ad = worker.active_done.clone();
            let transport_idle = transport_idle.clone();
            active_done = Some(worker.active_done);
            busy = Some(worker.busy.clone());
            std::thread::Builder::new()
                .name("pool-library-reader".into())
                .spawn(move || {
                    read_library(
                        worker.cmd_rx,
                        ad,
                        pp,
                        rr,
                        sr,
                        sd,
                        qf,
                        vcpu0_handle,
                        transport_idle,
                    )
                })
                .map_err(PoolError::ThreadSpawn)?;
        }

        Ok(Self {
            pause_pending,
            restore_req,
            snapshot_req,
            snapshot_done,
            quit,
            writer,
            supervisor_writer: Some(snapshot_done_writer),
            active_done,
            busy,
        })
    }

    /// Write `BAKE_READY` on the supervisor socket. Used by the
    /// runner when `dispatch_vcpu` returns `BakeReady` from the
    /// pipelined-bake flow. Best-effort; if there's no supervisor
    /// writer (in-process pool) or the writer's gone, silently no-
    /// op — nobody on the host is waiting for the line.
    pub fn signal_bake_ready(&self) {
        let Some(w) = self.supervisor_writer.as_ref() else {
            return;
        };
        if let Ok(mut g) = w.lock() {
            if let Some(s) = g.as_mut() {
                use std::io::Write;
                let _ = writeln!(s, "BAKE_READY");
                let _ = s.flush();
            }
        }
    }

    /// Clone of the supervisor writer Arc, for hand-off to bg save
    /// threads so they can post `SAVE_DONE <path>` when their write
    /// completes. Returns `None` for in-process pools (no supervisor
    /// to talk to). The Arc<Mutex<Option<UnixStream>>> contains the
    /// already-shared write half — bg threads serialize via the
    /// inner Mutex with other supervisor writers (DONE_SNAPSHOT,
    /// BAKE_READY, etc.), so `SAVE_DONE` can't interleave bytes
    /// mid-line with another message.
    pub fn supervisor_writer_clone(&self) -> Option<Arc<Mutex<Option<UnixStream>>>> {
        self.supervisor_writer.as_ref().map(Arc::clone)
    }

    pub fn should_quit(&self) -> bool {
        self.quit.load(Ordering::SeqCst)
    }

    pub fn pause_requested(&self) -> bool {
        self.pause_pending.load(Ordering::SeqCst)
    }

    pub fn pause_flag(&self) -> &AtomicBool {
        &self.pause_pending
    }

    pub fn clear_pause(&self) {
        self.pause_pending.store(false, Ordering::SeqCst);
    }

    /// Take any pending snapshot request the runner should
    /// process. Returned at most once per request — runner
    /// processes synchronously and posts the result via
    /// [`PoolControl::post_snapshot_result`].
    pub fn take_snapshot_request(&self) -> Option<SnapshotRequest> {
        self.snapshot_req.lock().ok()?.take()
    }

    /// Post the runner's snapshot result (Ok or Err) back to the
    /// host that called [`PoolHandle::snapshot`].
    pub fn post_snapshot_result(&self, result: Result<SnapshotResult, String>) {
        if let Ok(mut g) = self.snapshot_done.lock() {
            if let Some(tx) = g.take() {
                let _ = tx.send(result);
            }
        }
    }

    pub fn take_restore_request(&self) -> Option<RestoreRequest> {
        self.restore_req.lock().unwrap().take()
    }

    pub fn complete_restore(
        &mut self,
        us: u128,
        host_port: Option<u16>,
        timings: WarmRestoreTimings,
    ) {
        let result = PoolRestoreResult {
            restore_us: us,
            host_port,
            timings,
        };
        if let Some(w) = self.writer.as_mut() {
            let _ = write_done(w, us, host_port, timings);
        }
        if let Some(active_done) = self.active_done.as_ref() {
            if let Some(tx) = active_done.lock().unwrap().take() {
                let _ = tx.send(result);
            }
        }
        if let Some(busy) = self.busy.as_ref() {
            busy.store(false, Ordering::SeqCst);
        }
    }
}

fn write_done(
    w: &mut UnixStream,
    us: u128,
    host_port: Option<u16>,
    timings: WarmRestoreTimings,
) -> std::io::Result<()> {
    write!(w, "DONE {us}")?;
    if let Some(port) = host_port {
        write!(w, " host_port={port}")?;
    }
    writeln!(
        w,
        " reset_vsock_us={} remap_cow_us={} load_meta_us={} restore_snapshot_us={} ram_copy_us={} gic_restore_us={} vcpu_restore_us={} vtimer_offset_us={} mmio_restore_us={} listener_restore_us={}",
        timings.reset_vsock_us,
        timings.remap_cow_us,
        timings.load_meta_us,
        timings.restore_snapshot_us,
        timings.ram_copy_us,
        timings.gic_restore_us,
        timings.vcpu_restore_us,
        timings.vtimer_offset_us,
        timings.mmio_restore_us,
        timings.listener_restore_us
    )
}

#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn read_supervisor(
    read_half: UnixStream,
    pause_pending: Arc<AtomicBool>,
    restore_req: Arc<Mutex<Option<RestoreRequest>>>,
    snapshot_req: Arc<Mutex<Option<SnapshotRequest>>>,
    snapshot_done: Arc<Mutex<Option<mpsc::Sender<Result<SnapshotResult, String>>>>>,
    snapshot_done_writer: Arc<Mutex<Option<UnixStream>>>,
    quit: Arc<AtomicBool>,
    vcpu0_handle: crate::hypervisor::ActiveVcpuHandle,
    transport_idle: Option<TransportIdleCheck>,
    // This VM's per-VM PL011 state (`Vmm::serial`). Used by the
    // SMPARK_UNPARK_DIRECT bake fallback (reads `smpark_state_gpa`) and
    // PUSH_RX (calls `push_rx_byte`).
    serial: Arc<crate::devices::serial::SerialState>,
) {
    use std::io::BufRead;

    let mut reader = std::io::BufReader::new(read_half);
    // Per-worker smpark unpark-page GPA for the snapshot most recently
    // RESTOREd by this worker, read from that snapshot's `.smpark`
    // sidecar. Snapshot-scoped (correct across processes AND across
    // images sharing one host process) and loop-local (per-worker ⇒ no
    // cross-instance race) — unlike the process-global SMPARK_STATE_GPA
    // static, which holds the BOOTED guest's GPA and is stale/wrong for
    // a restored or different-image guest. Cached by path to avoid a
    // sidecar read every cycle.
    let mut cycle_smpark_gpa: u64 = 0;
    let mut cycle_smpark_path: Option<String> = None;
    loop {
        let mut line = String::new();
        match reader.read_line(&mut line) {
            Ok(0) | Err(_) => {
                quit.store(true, Ordering::SeqCst);
                pause_pending.store(true, Ordering::SeqCst);
                exit_vcpu(vcpu0_handle);
                break;
            }
            Ok(_) => {}
        }

        let cmd = line.trim();
        if let Some(rest) = cmd.strip_prefix("RESTORE ") {
            let mut parts = rest.split_ascii_whitespace();
            let path = parts.next().unwrap_or("").to_string();
            // Refresh the per-worker smpark GPA from THIS snapshot's
            // sidecar (cached by path). Consumed by SMPARK_UNPARK_DIRECT
            // below for host-direct multi-vCPU unpark.
            if cycle_smpark_path.as_deref() != Some(path.as_str()) {
                cycle_smpark_gpa = crate::vmm::snapshot::read_smpark_sidecar(&path).unwrap_or(0);
                cycle_smpark_path = Some(path.clone());
            }
            let mut egress_policy = None;
            for kv in parts {
                if let Some(v) = kv.strip_prefix("egress_policy=") {
                    egress_policy = Some(v.to_string());
                }
            }
            wait_until_transport_idle(transport_idle.as_ref(), &quit);
            *restore_req.lock().unwrap() = Some(RestoreRequest {
                path,
                egress_policy,
            });
            pause_pending.store(true, Ordering::SeqCst);
            exit_vcpu(vcpu0_handle);
        } else if let Some(rest) = cmd
            .strip_prefix("SNAPSHOT ")
            .or_else(|| cmd.strip_prefix("SNAPSHOT_ASYNC "))
        {
            // SNAPSHOT <out_path> (sync): pause guest, capture,
            // write to disk, post `DONE_SNAPSHOT …` back on the
            // supervisor socket.
            //
            // SNAPSHOT_ASYNC <out_path> (pipelined): pause guest
            // briefly, capture into a compact in-memory buffer,
            // spawn a background save thread, return
            // `DONE_SNAPSHOT_ASYNC` immediately. Used by the
            // bake-then-pool flow so the base snapshot save
            // overlaps with the warmup workload.
            let mode = if cmd.starts_with("SNAPSHOT_ASYNC ") {
                SnapshotMode::Async
            } else {
                SnapshotMode::Sync
            };
            let mut tokens = rest.split_ascii_whitespace();
            let out_path = tokens.next().unwrap_or("").to_string();
            // Trailing tokens. Currently:
            //   * `base=<path>` — diff-snapshot hint (Sync only).
            //   * `inline`      — wait for the diff save to land (Sync+base).
            // Unknown tokens are ignored for forward compatibility.
            let mut base_path = None;
            let mut inline_diff = false;
            for tok in tokens {
                if let Some(v) = tok.strip_prefix("base=") {
                    base_path = Some(v.to_string());
                } else if tok == "inline" {
                    inline_diff = true;
                }
            }
            let (done_tx, done_rx) = mpsc::channel();
            *snapshot_done.lock().unwrap() = Some(done_tx);
            wait_until_transport_idle(transport_idle.as_ref(), &quit);
            *snapshot_req.lock().unwrap() = Some(SnapshotRequest {
                out_path,
                mode,
                base_path,
                inline_diff,
            });
            pause_pending.store(true, Ordering::SeqCst);
            exit_vcpu(vcpu0_handle);
            // Wait for runner to post the result, then forward
            // it back to the lib side over the supervisor socket.
            // Bounded wait — if the runner crashed we don't want
            // this thread to hang forever; 60 s covers any
            // realistic disk-write of guest RAM.
            let resp = match done_rx.recv_timeout(Duration::from_secs(60)) {
                Ok(r) => r,
                Err(e) => Err(format!("snapshot wait: {e}")),
            };
            if let Ok(mut w) = snapshot_done_writer.lock() {
                if let Some(s) = w.as_mut() {
                    use std::io::Write;
                    match resp {
                        Ok(r) => {
                            let kind = if mode == SnapshotMode::Async {
                                "DONE_SNAPSHOT_ASYNC"
                            } else {
                                "DONE_SNAPSHOT"
                            };
                            let _ = writeln!(
                                s,
                                "{kind} bytes_written={} capture_us={} save_us={}",
                                r.bytes_written, r.capture_us, r.save_us
                            );
                        }
                        Err(e) => {
                            // Single line, replace newlines so the
                            // line-oriented protocol stays parseable.
                            let msg = e.replace('\n', " ");
                            let _ = writeln!(s, "ERR_SNAPSHOT {msg}");
                        }
                    }
                    let _ = s.flush();
                }
            }
        } else if cmd == "SMPARK_UNPARK_DIRECT" {
            // Host-direct smpark unpark. Bypasses the in-guest
            // agent vsock RPC entirely — the agent's
            // connection-handler thread can starve under heavy
            // multi-vCPU CPU contention, hitting the host-side
            // 5 s timeout. This path is a single memory store
            // from the host into the guest's `smpark_unpark_signal`
            // page; bounded by a syscall (~µs) and cannot be
            // CPU-starved by the guest.
            //
            // GPA discovery: smpark.ko publishes the page's GPA
            // via `/sys/kernel/smpark/state_gpa`; init-oci reads
            // that file after `finit_module` and prints it on
            // PL011 as `[SUPERMACHINE-SMPARK] state_gpa=0x…`;
            // the worker's PL011 scanner captures it into this VM's
            // `serial.smpark_state_gpa` (devices/serial.rs).
            //
            // Older smpark.ko without the sysfs file → GPA is 0
            // → return ERR_SMPARK_UNPARK; the bake driver falls
            // back to the agent vsock RPC path.
            //
            // The parked secondaries wake on the next vtimer
            // tick (≤10 ms) — WFI returns on ANY pending IRQ
            // even when CPSR.I is masked, so no IPI is needed.
            use std::sync::atomic::Ordering;
            // Prefer the restored snapshot's `.smpark` sidecar GPA
            // (per-worker, set in the RESTORE handler above): it is
            // snapshot-scoped, so a RESTORED guest unparks with the GPA
            // ITS snapshot was baked with. When it's 0 — i.e. no RESTORE
            // happened in this worker, which is the BAKE worker's
            // post-capture unpark of a freshly-BOOTED guest — fall back
            // to this VM's `serial.smpark_state_gpa`, which the PL011
            // boot-line scanner just set to THIS guest's GPA. Without
            // this fallback the bake's host-direct unpark always
            // reported "not announced" and fell back to the agent RPC,
            // which can starve under multi-vCPU + heavy-warmup
            // contention and hit the 5 s timeout. (The serial state is
            // per-VM since 0.7.56, so there is no longer any cross-image
            // staleness to guard against — but the sidecar GPA still
            // wins for restores because the booted-guest scanner value
            // is only relevant to the bake worker's own boot.) 0 from
            // both ⇒ ERR ⇒ agent-RPC fallback below.
            let gpa = if cycle_smpark_gpa != 0 {
                cycle_smpark_gpa
            } else {
                serial.smpark_state_gpa.load(Ordering::Acquire)
            };
            let ack: Result<(), String> = if gpa == 0 {
                Err(
                    "smpark state_gpa not announced — older smpark.ko, or smpark.ko absent"
                        .to_owned(),
                )
            } else {
                let base = crate::vmm::vstate::HOST_RAM_GPA.load(Ordering::Acquire);
                let size = crate::vmm::vstate::HOST_RAM_SIZE.load(Ordering::Acquire);
                let ptr = crate::vmm::vstate::HOST_RAM_PTR.load(Ordering::Acquire);
                if ptr.is_null() {
                    Err("HOST_RAM_PTR not published — runner not ready".to_owned())
                } else if gpa < base || gpa >= base + size {
                    Err(format!(
                        "GPA 0x{gpa:x} outside RAM range [0x{base:x}, 0x{:x})",
                        base + size
                    ))
                } else {
                    let off = (gpa - base) as usize;
                    // SAFETY: ptr is the published RAM host-VA;
                    // offset is bounds-checked against `size`;
                    // single-byte volatile write. The parked
                    // secondary's WFI loop ends in a `READ_ONCE`
                    // which compiles to a plain `LDR` — adequate
                    // because WFI itself acts as a barrier when
                    // it exits.
                    unsafe {
                        let hva = ptr.add(off);
                        std::ptr::write_volatile(hva, 1u8);
                    }
                    Ok(())
                }
            };
            if let Ok(mut w) = snapshot_done_writer.lock() {
                if let Some(s) = w.as_mut() {
                    use std::io::Write;
                    match ack {
                        Ok(_) => {
                            let _ = writeln!(s, "DONE_SMPARK_UNPARK");
                        }
                        Err(e) => {
                            let msg = e.replace('\n', " ");
                            let _ = writeln!(s, "ERR_SMPARK_UNPARK {msg}");
                        }
                    }
                    let _ = s.flush();
                }
            }
        } else if let Some(rest) = cmd.strip_prefix("PUSH_RX ") {
            // 0.7.43+ kernel-boot cache resume protocol. The bake
            // driver sends `PUSH_RX <hex_byte>` to deliver a byte
            // to the guest via PL011 RX. Multiple PUSH_RX in
            // sequence build up the resume token (typically "R\n"
            // → "PUSH_RX 52" + "PUSH_RX 0a").
            //
            // No supervisor-side acknowledgement — the bake driver
            // observes resume by reading the next line on the
            // worker's stderr (init-oci's
            // "[SUPERMACHINE-KCACHE] resumed token=…").
            let trimmed = rest.trim();
            if let Ok(b) = u8::from_str_radix(trimmed, 16) {
                serial.push_rx_byte(b);
            } else {
                eprintln!("PUSH_RX: bad hex byte {trimmed:?}");
            }
        } else if cmd == "QUIT" {
            quit.store(true, Ordering::SeqCst);
            pause_pending.store(true, Ordering::SeqCst);
            exit_vcpu(vcpu0_handle);
            break;
        }
    }
}

#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn read_library(
    cmd_rx: mpsc::Receiver<PoolCommand>,
    active_done: Arc<Mutex<Option<mpsc::Sender<PoolRestoreResult>>>>,
    pause_pending: Arc<AtomicBool>,
    restore_req: Arc<Mutex<Option<RestoreRequest>>>,
    snapshot_req: Arc<Mutex<Option<SnapshotRequest>>>,
    snapshot_done: Arc<Mutex<Option<mpsc::Sender<Result<SnapshotResult, String>>>>>,
    quit: Arc<AtomicBool>,
    vcpu0_handle: crate::hypervisor::ActiveVcpuHandle,
    transport_idle: Option<TransportIdleCheck>,
) {
    loop {
        match cmd_rx.recv() {
            Ok(PoolCommand::Restore { request, done_tx }) => {
                *active_done.lock().unwrap() = Some(done_tx);
                wait_until_transport_idle(transport_idle.as_ref(), &quit);
                *restore_req.lock().unwrap() = Some(request);
                pause_pending.store(true, Ordering::SeqCst);
                exit_vcpu(vcpu0_handle);
            }
            Ok(PoolCommand::Snapshot { request, done_tx }) => {
                // Same shape as Restore: stash the request +
                // completion channel, force the vCPU out of
                // dispatch, runner picks up the request in its
                // post-dispatch loop.
                *snapshot_done.lock().unwrap() = Some(done_tx);
                wait_until_transport_idle(transport_idle.as_ref(), &quit);
                *snapshot_req.lock().unwrap() = Some(request);
                pause_pending.store(true, Ordering::SeqCst);
                exit_vcpu(vcpu0_handle);
            }
            Ok(PoolCommand::Quit) => {
                active_done.lock().unwrap().take();
                quit.store(true, Ordering::SeqCst);
                pause_pending.store(true, Ordering::SeqCst);
                exit_vcpu(vcpu0_handle);
                break;
            }
            Err(_) => {
                active_done.lock().unwrap().take();
                quit.store(true, Ordering::SeqCst);
                pause_pending.store(true, Ordering::SeqCst);
                exit_vcpu(vcpu0_handle);
                break;
            }
        }
    }
}

#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn wait_until_transport_idle(transport_idle: Option<&TransportIdleCheck>, quit: &AtomicBool) {
    let Some(transport_idle) = transport_idle else {
        return;
    };
    // Cap the wait at SUPERMACHINE_TRANSPORT_IDLE_WAIT_MS
    // (default 5 000 ms). Field-report 0.7.32: a warmup that
    // leaves background TCP/UDP activity going (e.g. `npm
    // install` postinstall hooks doing CDN fetches) keeps the
    // TSI muxer's `pending_*` / `inbound_conns` / `rxq` non-
    // empty indefinitely — the prior unbounded spin loop here
    // wedged the supervisor BEFORE it could even send the
    // snapshot request to the runner, surfacing as a worker
    // burning CPU forever post-`warmup done` with no further
    // log output.
    //
    // With a cap, the snapshot proceeds with whatever transport
    // state is in flight. Consequences:
    //   * Warm-handoff (same-process first acquire) is
    //     unaffected — the live muxer state continues running
    //     normally alongside the snapshot capture; in-flight
    //     connections aren't "lost" because they were never
    //     captured (they keep running in the same process).
    //   * Pool spawn-from-disk on the resulting snapshot loses
    //     in-flight connections, which is the same outcome any
    //     disk-restore has — workloads relying on stable cross-
    //     restore TCP state aren't a supported pattern.
    //
    // 5 s is enough that lightly-busy muxers (occasional
    // background pings) typically drain naturally; persistently
    // chatty muxers (the integrator's case) fall through
    // promptly so the bake can finish. Override via
    // `SUPERMACHINE_TRANSPORT_IDLE_WAIT_MS=N` (any integer
    // milliseconds) — set to `0` to skip the wait entirely;
    // set to a large value like `60000` to restore near-
    // unbounded waiting for workloads that genuinely settle
    // within a minute but not within 5 s.
    let cap_ms: u64 = std::env::var("SUPERMACHINE_TRANSPORT_IDLE_WAIT_MS")
        .ok()
        .and_then(|s| s.parse().ok())
        .unwrap_or(5_000);
    let deadline = std::time::Instant::now() + Duration::from_millis(cap_ms);
    while !quit.load(Ordering::SeqCst) && !transport_idle() {
        if std::time::Instant::now() > deadline {
            eprintln!(
                "  [pool] WARNING: wait_until_transport_idle hit the {cap_ms} ms \
                 cap with the TSI muxer still busy. The snapshot will be \
                 captured with in-flight connections — clones restored from it \
                 (disk / cross-process) see those as half-open and may hang \
                 until the app times out. Quiesce network activity at the end \
                 of your warmup, or raise SUPERMACHINE_TRANSPORT_IDLE_WAIT_MS, \
                 to avoid baking a mid-transfer snapshot.",
            );
            return;
        }
        std::thread::sleep(Duration::from_micros(100));
    }
}

#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn exit_vcpu(vcpu0_handle: crate::hypervisor::ActiveVcpuHandle) {
    // Force vCPU 0 out of its run loop via the backend-agnostic seam (HVF:
    // `hv_vcpus_exit`; KVM: signal the vCPU thread). The handle is valid for the
    // lifetime of the pool reader.
    crate::hypervisor::VcpuHandle::force_exit(std::slice::from_ref(&vcpu0_handle));
}
