nucleic/Sources/NucleicCore/MacVM/MacVMEngine+Provision.swift

import Foundation
#if arch(arm64)
import Virtualization
#endif

/// **Toolchain + native-agent provisioning of the golden base, driven over SSH** — the second half of
/// the one-click "Build base image" flow (docs/MACOS_VM.md §4). After the install +
/// account provisioning (`MacVMEngine+Provision27.swift`) leaves a key-authorized `agent` account,
/// this pass boots the base once more, stages `scripts/provision-macos-guest.sh` (+ the host public
/// key + the native `NucleicVMAgent.app`) over a virtiofs share, and runs the provisioner
/// unattended: dev toolchain (Xcode CLT, Homebrew, node/python, the `/etc/zshenv` PATH), then — when
/// computer use is enabled — the agent app + LaunchAgent + its computer-use TCC grants.
///
/// The provisioner's TCC writes need **SIP disabled** in the guest (docs/MACOS_VM.md §12.5), which
/// can't be scripted over SSH; Phase 7 detects SIP and skips those grants gracefully when it's on.
/// So this pass always makes the base **exec-ready** and, when SIP is already off, **computer-use-
/// ready** too; otherwise it reports `computerUseReady == false` and the UI surfaces the one-time
/// Recovery step (`bootBaseInRecovery`), after which a second Build re-runs this pass and the grants
/// land.
extension MacVMEngine {
    /// What the provisioning pass achieved, folded into the base's `bundle.json`.
    struct ProvisionResult: Sendable {
        var provisioned: Bool        // the toolchain script ran to completion (exec-ready)
        var agentInstalled: Bool     // NucleicVMAgent.app + LaunchAgent present in the guest
        var sipDisabled: Bool        // SIP off (⇒ the AX agent's TCC grants could be written)
        var axAgentReady: Bool       // AX agent installed AND its TCC grants are in place
    }

    /// Parsed readback of the `CODE=… SIP=… APP=… LA=… TCC=…` line the guest bootstrap writes to the
    /// STATUS file on the shared directory after provisioning (docs/MACOS_VM.md §4.4).
    struct ProvisionStatusFlags: Sendable, Equatable {
        var exitCode: Int32 = -1
        var sipDisabled = false
        var agentApp = false
        var launchAgent = false
        var tccCount = 0
    }

    /// Parse the `CODE=0 SIP=1 APP=1 LA=1 TCC=3` readback line (order-independent, tolerant of noise).
    static func parseProvisionStatus(_ output: String) -> ProvisionStatusFlags {
        var flags = ProvisionStatusFlags()
        for token in output.split(whereSeparator: { $0 == " " || $0.isNewline }) {
            let parts = token.split(separator: "=", maxSplits: 1)
            guard parts.count == 2 else { continue }
            let value = String(parts[1]).trimmingCharacters(in: .whitespaces)
            switch parts[0] {
            case "CODE": flags.exitCode = Int32(value) ?? -1
            case "SIP": flags.sipDisabled = (value == "1")
            case "APP": flags.agentApp = (value == "1")
            case "LA": flags.launchAgent = (value == "1")
            case "TCC": flags.tccCount = Int(value) ?? 0
            default: break
            }
        }
        return flags
    }

    // MARK: - Provisioning-asset resolution (bundled in the app; dev fallback to the source tree)

    /// The staged `scripts/provision-macos-guest.sh` — bundled into the app under `Resources/macvm/`
    /// (`scripts/package-app.sh`), or, for a dev/SwiftPM run, the copy in the source tree.
    func resolveProvisionScript() -> URL? {
        if let bundled = Bundle.main.resourceURL?
            .appendingPathComponent("macvm/provision-macos-guest.sh"),
            FileManager.default.fileExists(atPath: bundled.path)
        {
            return bundled
        }
        return Self.devRepoRoot()?.appendingPathComponent("scripts/provision-macos-guest.sh")
    }

    /// The signed `NucleicVMAgent.app` to bake in — bundled under `Resources/macvm/`, or, for a dev
    /// run, the `scripts/build-vm-agent.sh` output (built on demand if absent). `nil` when neither
    /// exists (computer use is then unavailable and the toolchain-only base is still exec-ready).
    func resolveVMAgentApp() -> URL? {
        if let bundled = Bundle.main.resourceURL?
            .appendingPathComponent("macvm/NucleicVMAgent.app"),
            FileManager.default.fileExists(atPath: bundled.path)
        {
            return bundled
        }
        guard let root = Self.devRepoRoot() else { return nil }
        let dist = root.appendingPathComponent("guest/NucleicVMAgent/dist/NucleicVMAgent.app")
        if FileManager.default.fileExists(atPath: dist.path) { return dist }
        // Dev convenience: build it once via the same script the packaging step uses.
        let build = Process()
        build.executableURL = URL(fileURLWithPath: "/bin/bash")
        build.arguments = [root.appendingPathComponent("scripts/build-vm-agent.sh").path]
        build.standardOutput = Pipe()
        build.standardError = Pipe()
        try? build.run()
        build.waitUntilExit()
        return FileManager.default.fileExists(atPath: dist.path) ? dist : nil
    }

    /// The LaunchAgent plist that rides alongside the agent app (same directory as `resolveVMAgentApp`).
    func resolveVMAgentPlist(near app: URL) -> URL? {
        let plist = app.deletingLastPathComponent()
            .appendingPathComponent("xyz.blakeslee.nucleic.vmagent.plist")
        return FileManager.default.fileExists(atPath: plist.path) ? plist : nil
    }

    /// Locate the repo root from this source file (dev/SwiftPM runs only): ascend until
    /// `scripts/provision-macos-guest.sh` is found. `nil` from a shipped app (no source tree).
    static func devRepoRoot() -> URL? {
        var dir = URL(fileURLWithPath: #filePath).deletingLastPathComponent()
        for _ in 0..<6 {
            if FileManager.default.fileExists(
                atPath: dir.appendingPathComponent("scripts/provision-macos-guest.sh").path)
            {
                return dir
            }
            dir = dir.deletingLastPathComponent()
        }
        return nil
    }
}

#if arch(arm64)
extension MacVMEngine {
    /// Run the account + toolchain + agent provisioning pass on a freshly installed (or already
    /// account-provisioned) base and fold the result into its `bundle.json`. Best-effort on *content*
    /// — a SIP-on guest yields an exec-ready-but-not-AX base — and throws only when the guest never
    /// became drivable, leaving the published base intact for a retry (the reentrant "Build again").
    ///
    /// `declarativeFirstBoot` (macOS 27 fresh install) makes this the first boot after restore, so it
    /// carries `VZMacGuestProvisioningOptions` to create the `agent` account + auto-login.
    func provisionAndFinalize(
        bundle: MacVMBundle, base: MacVMBaseStatus, declarativeFirstBoot: Bool
    ) async throws {
        baseProgress = MacVMBaseProgress(
            phase: declarativeFirstBoot ? .firstBootSetup : .provisioning, fraction: nil)
        do {
            // The native agent is the sole host↔guest control channel now, so it is ALWAYS installed
            // (exec + computer-use ride it). Its AX *TCC grants* remain SIP-gated and are the only part
            // `MacVMSettings.axAgentEnabled` affects downstream (see `axAgentReady`).
            let result = try await provisionBase(
                bundle: bundle, installAgent: true, declarativeFirstBoot: declarativeFirstBoot)
            var status = base
            // Reaching a shell + running the bootstrap means the account exists and auto-login works.
            status.accountProvisioned = status.accountProvisioned || declarativeFirstBoot
            status.provisioned = result.provisioned
            status.agentInstalled = result.agentInstalled
            status.sipDisabled = result.sipDisabled
            status.axAgentReady = result.axAgentReady
            writeBaseStatus(bundle, status)
            baseProgress = MacVMBaseProgress(phase: .ready, fraction: 1)
        } catch {
            baseProgress = nil
            throw MacVMError.provisionFailed(String(describing: error))
        }
    }

    /// Provision the base **entirely host-side + VirtioFS, no network** (docs/MACOS_VM.md §4.4):
    ///   1. stage the provisioner, the signed agent app, the account password, and a host-composed
    ///      bootstrap into ONE read-write virtiofs share;
    ///   2. boot the base (declaratively creating the account on a fresh macOS-27 install) with a
    ///      host-side `VZVirtualMachineView` surface bound;
    ///   3. drive the surface's synthesized keyboard to open Terminal and launch the staged bootstrap
    ///      — which primes sudo, runs `provision-macos-guest.sh` (dev toolchain + the native agent +,
    ///      when SIP is off, its TCC grants), writes a `STATUS` readback to the share, and powers off;
    ///   4. poll the share for `STATUS` (which also proves the desktop came up), then await shutdown.
    ///
    /// Throws only when the guest never became drivable (no `STATUS` before the deadline); a run that
    /// completes with partial results (SIP on ⇒ no TCC grants) returns partial flags, so the base is
    /// still published exec-ready.
    func provisionBase(
        bundle: MacVMBundle, installAgent: Bool, declarativeFirstBoot: Bool
    ) async throws -> ProvisionResult {
        let user = MacVMSettings.sshUser
        let password = try resolveAgentPassword()
        guard let scriptURL = resolveProvisionScript() else {
            throw MacVMError.provisionFailed(
                "the bundled provision-macos-guest.sh could not be located")
        }
        guard let surface = surfaceHost else {
            throw MacVMError.provisionFailed(
                "base provisioning drives the guest through the host-side display surface, which "
                    + "isn't available in this context (headless/spike). Build the base from the app.")
        }

        // ── 1. Stage everything into one READ-WRITE virtiofs share (the guest writes STATUS back).
        let fm = FileManager.default
        let stageDir = fm.temporaryDirectory
            .appendingPathComponent("nucleic-provision-\(UUID().uuidString)", isDirectory: true)
        try fm.createDirectory(at: stageDir, withIntermediateDirectories: true)
        defer { try? fm.removeItem(at: stageDir) }

        try? fm.copyItem(
            at: scriptURL, to: stageDir.appendingPathComponent("provision-macos-guest.sh"))
        try password.write(
            to: stageDir.appendingPathComponent("password"), atomically: true, encoding: .utf8)

        var agentStaged = false
        if installAgent, let app = resolveVMAgentApp() {
            try? fm.copyItem(at: app, to: stageDir.appendingPathComponent("NucleicVMAgent.app"))
            if let plist = resolveVMAgentPlist(near: app) {
                try? fm.copyItem(
                    at: plist, to: stageDir.appendingPathComponent(plist.lastPathComponent))
            }
            agentStaged = fm.fileExists(
                atPath: stageDir.appendingPathComponent("NucleicVMAgent.app").path)
        }

        let bootstrapURL = stageDir.appendingPathComponent(Self.bootstrapScriptName)
        try Self.provisionBootstrapScript(user: user).write(
            to: bootstrapURL, atomically: true, encoding: .utf8)
        try? fm.setAttributes([.posixPermissions: 0o755], ofItemAtPath: bootstrapURL.path)

        // ── 2. Boot with the RW share mounted, on the MAIN queue so a host surface can bind.
        let mac = Self.randomMAC()
        let mount = MacVMSpec.Mount(host: stageDir.path, name: "nucleic-provision", readOnly: false)
        let config = try Self.makeConfiguration(
            bundle: bundle, cpus: MacVMSettings.vmCPUs, memoryGiB: MacVMSettings.vmMemoryGiB,
            mac: mac, mounts: [mount])
        let instance = MacVMInstance(
            configuration: config, label: Self.baseProvisionSurfaceName, mainQueue: true)
        do {
            if declarativeFirstBoot, #available(macOS 27.0, *) {
                try await instance.startWithProvisioning(
                    fullName: "Nucleic Agent", username: user, password: password)
            } else {
                try await instance.start()
            }
        } catch {
            return ProvisionResult(
                provisioned: false, agentInstalled: false, sipDisabled: false, axAgentReady: false)
        }

        let surfaceName = "base-provision"
        await surface.attach(
            name: surfaceName,
            virtualMachine: UncheckedSendableBox(value: instance.vm as AnyObject))

        // ── 3–4. Drive HID → launch the bootstrap → poll STATUS on the share.
        let activePhase: MacVMBaseProgress.Phase =
            declarativeFirstBoot ? .firstBootSetup : .installingAgent
        baseProgress = MacVMBaseProgress(phase: activePhase, fraction: nil)
        let statusURL = stageDir.appendingPathComponent("STATUS")
        let startedURL = stageDir.appendingPathComponent("STARTED")
        let logURL = stageDir.appendingPathComponent("provision.log")
        let flags = await driveHIDBootstrap(
            surface: surface, name: surfaceName, phase: activePhase,
            startedURL: startedURL, statusURL: statusURL, logURL: logURL)

        baseProgress = MacVMBaseProgress(phase: .finalizing, fraction: nil)
        await surface.detach(name: surfaceName)
        await instance.awaitStopped(timeout: 120)  // the bootstrap powers off when done

        guard let flags else {
            throw MacVMError.agentUnavailable(
                "the base never reached a usable desktop during provisioning (no STATUS readback)")
        }
        let provisioned = flags.exitCode == 0
        let agentInstalled = agentStaged && flags.agentApp && flags.launchAgent
        let axAgentReady = agentInstalled && flags.sipDisabled && flags.tccCount >= 3
        return ProvisionResult(
            provisioned: provisioned, agentInstalled: agentInstalled,
            sipDisabled: flags.sipDisabled, axAgentReady: axAgentReady)
    }

    /// The staged bootstrap's filename on the share.
    static let bootstrapScriptName = "nucleic-bootstrap.sh"

    /// Drive the host-side surface to launch the staged bootstrap, then poll the share for its
    /// `STATUS` readback. Re-attempts the launch (open Spotlight → Terminal → run) until a `STARTED`
    /// sentinel proves the bootstrap is running (so we stop typing into a live install), then just
    /// waits for `STATUS`. Returns `nil` if neither appears before the deadline.
    ///
    /// Each poll it republishes ``baseProgress`` with a live `detail`: before `STARTED`, that we're
    /// still trying to reach the desktop and launch the provisioner (the blind-HID step is the usual
    /// place a build wedges); after, the tail of the guest's `provision.log`, so the long, fraction-less
    /// toolchain install shows the step it's actually on instead of a frozen spinner.
    private func driveHIDBootstrap(
        surface: any MacVMSurfaceHost, name: String, phase: MacVMBaseProgress.Phase,
        startedURL: URL, statusURL: URL, logURL: URL
    ) async -> ProvisionStatusFlags? {
        let fm = FileManager.default
        let deadline = Date().addingTimeInterval(1800)  // ≤30 min: first login + toolchain install
        var running = false
        var nextLaunch = Date()
        while Date() < deadline {
            if let text = try? String(contentsOf: statusURL, encoding: .utf8), text.contains("CODE=") {
                await surface.setKeyboardFocus(name: name, focused: false)
                return Self.parseProvisionStatus(text)
            }
            if !running, fm.fileExists(atPath: startedURL.path) {
                running = true  // the bootstrap is executing — stop typing, hand the keyboard back
                await surface.setKeyboardFocus(name: name, focused: false)
            }
            if !running, Date() >= nextLaunch {
                // VZ only forwards keys from a key window, so take key for the duration of the typing.
                await surface.setKeyboardFocus(name: name, focused: true)
                await Self.typeBootstrapLaunch(surface: surface, name: name)
                nextLaunch = Date().addingTimeInterval(40)  // retry login-timing until STARTED appears
            }
            let detail = running
                ? (Self.lastProvisionLogLine(at: logURL) ?? "Running the provisioner in the guest…")
                : "Waiting for the guest desktop, then launching the provisioner…"
            baseProgress = MacVMBaseProgress(phase: phase, fraction: nil, detail: detail)
            try? await Task.sleep(nanoseconds: 3_000_000_000)
        }
        await surface.setKeyboardFocus(name: name, focused: false)  // released on timeout too
        return nil
    }

    /// The last non-empty line of the guest's `provision.log`, trimmed and length-bounded — the live
    /// step to surface as base-build progress while the (indeterminate) toolchain install runs. `nil`
    /// when the log isn't readable yet.
    static func lastProvisionLogLine(at url: URL) -> String? {
        guard let text = try? String(contentsOf: url, encoding: .utf8) else { return nil }
        guard let line = text.split(whereSeparator: \.isNewline)
            .last(where: { !$0.trimmingCharacters(in: .whitespaces).isEmpty })
        else { return nil }
        let trimmed = line.trimmingCharacters(in: .whitespaces)
        return trimmed.count > 140 ? String(trimmed.prefix(140)) + "…" : trimmed
    }

    /// Synthesize the keystrokes that open Terminal via Spotlight and run the staged bootstrap. Timed
    /// generously — this is a one-time base build, not a latency-sensitive path.
    private static func typeBootstrapLaunch(surface: any MacVMSurfaceHost, name: String) async {
        func pause(_ seconds: Double) async {
            try? await Task.sleep(nanoseconds: UInt64(seconds * 1_000_000_000))
        }
        // Dismiss anything focus-stealing, open Spotlight, search + launch Terminal.
        await surface.send(name: name, .key(chord: "esc"))
        await pause(0.5)
        await surface.send(name: name, .key(chord: "cmd+space"))
        await pause(1.0)
        await surface.send(name: name, .text("Terminal"))
        await pause(1.0)
        await surface.send(name: name, .key(chord: "return"))
        await pause(3.0)  // Terminal cold-launch
        // Run the staged bootstrap from the share.
        let command = "/bin/bash '/Volumes/My Shared Files/nucleic-provision/\(bootstrapScriptName)'"
        await surface.send(name: name, .text(command))
        await surface.send(name: name, .key(chord: "return"))
    }

    /// The host-composed bootstrap the guest runs (via HID) — network-free: it primes passwordless
    /// sudo, runs the provisioner, writes a `CODE=… SIP=… APP=… LA=… TCC=…` readback to `STATUS` on
    /// the share, removes the sudo drop-in, and powers off. Reads the account password from the staged
    /// `password` file (no password typing). The `STARTED` sentinel tells the host to stop re-typing.
    static func provisionBootstrapScript(user: String) -> String {
        let plist = "/Library/LaunchAgents/xyz.blakeslee.nucleic.vmagent.plist"
        let tccDB = "/Library/Application Support/com.apple.TCC/TCC.db"
        let tccQuery =
            "SELECT count(*) FROM access WHERE service IN "
            + "('kTCCServiceScreenCapture','kTCCServiceAccessibility','kTCCServicePostEvent') "
            + "AND auth_value=2;"
        return """
        #!/bin/bash
        SHARE="/Volumes/My Shared Files/nucleic-provision"
        touch "$SHARE/STARTED" 2>/dev/null || true
        PW="$(cat "$SHARE/password")"
        # Prime passwordless sudo so the multi-minute toolchain install can't stall on a tty/timestamp.
        printf '%s\\n' "$PW" | sudo -S sh -c 'grep -q "includedir /private/etc/sudoers.d" /etc/sudoers || printf "%s\\n" "@includedir /private/etc/sudoers.d" >> /etc/sudoers; printf "%s\\n" "\(user) ALL=(ALL) NOPASSWD: ALL" > /etc/sudoers.d/nucleic-provision; chmod 0440 /etc/sudoers.d/nucleic-provision' 2>/dev/null
        NUCLEIC_AGENT_PW="$PW" NUCLEIC_PROVISION_NO_SHUTDOWN=1 /bin/bash "$SHARE/provision-macos-guest.sh" "" "$PW" > "$SHARE/provision.log" 2>&1
        CODE=$?
        SIP=$(csrutil status 2>/dev/null | grep -qi disabled && echo 1 || echo 0)
        APP=$(test -d /Applications/NucleicVMAgent.app && echo 1 || echo 0)
        LA=$(test -f "\(plist)" && echo 1 || echo 0)
        TCC=$(sudo sqlite3 "\(tccDB)" "\(tccQuery)" 2>/dev/null || echo 0)
        printf 'CODE=%s SIP=%s APP=%s LA=%s TCC=%s\\n' "$CODE" "$SIP" "$APP" "$LA" "$TCC" > "$SHARE/STATUS"
        sync
        sudo rm -f /etc/sudoers.d/nucleic-provision
        sudo shutdown -h now
        """
    }
}
#endif