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

/// **Install "Common packages" (``MacVMPackage``) into a guest over vsock** — the runtime counterpart
/// to baking them into the base at build time (`scripts/provision-macos-guest.sh`, via the staged
/// `install-packages.sh` the engine writes in `MacVMEngine+Provision.swift`). Two entry points share
/// one core, mirroring `MacVMEngine+InstallApps.swift`:
///
///  * ``installPackages(name:packageIDs:)`` targets a RUNNING session clone (Settings → "Copy to
///    running VMs"), so a live VM gets the packages without a restart.
///  * ``installPackagesIntoBase(packageIDs:)`` briefly boots the **golden base** writable, installs the
///    packages, and powers it back off — so newly-added packages persist into every FUTURE clone
///    WITHOUT the ~30-minute rebuild (Settings → "Add to base image").
///
/// Unlike apps, there's nothing to stream from the host: each package's install shell
/// (``MacVMPackage/installScript``) fetches what it needs inside the guest (e.g. Chrome's public DMG),
/// so both paths just run ``MacVMPackage/installShell(forIDs:)`` over the native agent's `exec`.
extension MacVMEngine {
    /// Install each installable package in `packageIDs` into the named **running** guest. Returns the
    /// package ids that actually landed (parsed from the run's success markers). Throws
    /// ``MacVMError/notRunning(_:)`` if the VM isn't live, or ``MacVMError/appInstallFailed(_:)`` if the
    /// guest can't be reached. An empty / all-unknown set is a no-op returning `[]`.
    @discardableResult
    public func installPackages(name: String, packageIDs: [String]) async throws -> [String] {
        #if arch(arm64)
        guard let instance = live[name]?.instance else { throw MacVMError.notRunning(name) }
        return try await installPackages(onInstance: instance, label: name, packageIDs: packageIDs)
        #else
        throw MacVMError.unavailable(Self.unsupportedReason ?? "requires Apple silicon")
        #endif
    }

    /// Install the given packages into the **golden base image** without a full rebuild: boot the base
    /// writable + headless (its baked-in agent auto-logs-in and listens on vsock), run the installer,
    /// then shut it down cleanly so the quiesced disk can be cloned again. New session VMs cut from the
    /// base afterward include the packages. Returns the ids that landed.
    ///
    /// Refuses (``MacVMError/baseBusy(_:)``) while the base is otherwise in use, throws
    /// ``MacVMError/baseImageMissing(_:)`` when no base is built, and ``MacVMError/appInstallFailed(_:)``
    /// if the base can't be booted/reached (e.g. the OS's concurrent-guest ceiling is already reached).
    @discardableResult
    public func installPackagesIntoBase(packageIDs: [String]) async throws -> [String] {
        #if arch(arm64)
        guard MacVMPackage.installShell(forIDs: packageIDs) != nil else { return [] }
        guard !baseIsBusy else {
            throw MacVMError.baseBusy(
                baseBuilding ? "a base image build is in progress"
                    : baseRecoveryActive ? "the base is open in a Recovery window"
                    : "the base is already installing apps")
        }
        let base = try await ensureBaseBundle()  // throws .baseImageMissing when none exists
        guard readBaseStatus(base).accountProvisioned else {
            throw MacVMError.appInstallFailed(
                "the base image isn't provisioned yet (it has no in-guest agent to reach). Build the "
                    + "base image first, then add packages to it.")
        }

        // The OS caps concurrent macOS guests; a base boot is one more. Surface a friendly message
        // instead of an opaque VZ start failure when sessions already fill the ceiling.
        guard live.count < MacVMSettings.maxConcurrentVMs else {
            throw MacVMError.appInstallFailed(
                "\(live.count) macOS VM\(live.count == 1 ? " is" : "s are") already running (the OS "
                    + "limits concurrent guests). Stop one, or add the packages to the running VMs instead.")
        }

        baseAppInstalling = true
        // Publish the base as a running, system-managed VM so the Control panel shows it (with a live
        // activity note) alongside the per-session guests while packages install.
        baseMaintenance = MacVMMaintenanceInfo(
            name: "base-pkg-install", os: .macOS, activity: "Starting…", booting: true)
        defer {
            baseAppInstalling = false
            baseMaintenance = nil
        }

        let config = try Self.makeConfiguration(
            bundle: base, cpus: MacVMSettings.vmCPUs, memoryGiB: MacVMSettings.vmMemoryGiB,
            mac: Self.randomMAC(), mounts: [])
        let instance = MacVMInstance(configuration: config, label: "base-pkg-install")
        do {
            try await instance.start()
        } catch {
            throw MacVMError.appInstallFailed("could not boot the base image: \(error)")
        }
        do {
            try await awaitAgentReady(onInstance: instance)
            baseMaintenance = MacVMMaintenanceInfo(
                name: "base-pkg-install", os: .macOS, activity: "Installing packages…", booting: false)
            // Bound the install so a wedged guest can't hang the operation — and leave the base booted —
            // forever (a package download plus the wait for the guest's `exit` frame).
            let installed: [String]
            do {
                installed = try await withAgentTimeout(seconds: 2000) { [packageIDs] in
                    try await self.installPackages(
                        onInstance: instance, label: "base", packageIDs: packageIDs)
                }
            } catch is MacVMAgentError {  // only the timeout injects this; installPackages throws MacVMError
                throw MacVMError.appInstallFailed(
                    "the base image stopped responding while installing packages (timed out). It has "
                        + "been powered off — try again.")
            }
            await shutdownBaseGuest(instance)  // clean power-off so the disk quiesces before re-cloning
            return installed
        } catch {
            await instance.stop()  // force power-off on any failure; the base disk is unchanged on error
            throw error
        }
        #else
        throw MacVMError.unavailable(Self.unsupportedReason ?? "requires Apple silicon")
        #endif
    }

    #if arch(arm64)
    /// Shared core: run ``MacVMPackage/installShell(forIDs:)`` in `instance` over `exec` and return the
    /// package ids the run reported installing (from its success markers). `label` names the target only
    /// for error text. A set with nothing installable is a no-op returning `[]`.
    private func installPackages(
        onInstance instance: MacVMInstance, label: String, packageIDs: [String]
    ) async throws -> [String] {
        guard let body = MacVMPackage.installShell(forIDs: packageIDs) else { return [] }

        let channel: MacVMExecChannel
        do {
            channel = try await openExecChannel(
                instance: instance, label: label, workdir: nil, env: [:], remoteBody: body)
        } catch {
            throw MacVMError.appInstallFailed("\(error)")
        }
        channel.closeStdin()  // the installer reads no stdin; signal EOF so the guest shell can finish

        // Drain both streams concurrently so the exit frame is processed and stdout (with the success
        // markers) is captured. Per-package failures are logged to stderr but not fatal, so the exit
        // code reflects the shell as a whole; which packages landed is read from the markers.
        async let stderrText = Self.drainLines(channel.stderrLines)
        let stdoutText = await Self.drainLines(channel.stdoutLines)
        let code = await channel.wait()
        let stderr = await stderrText

        guard code == 0 else {
            let detail = stderr.isEmpty ? "guest exit \(code)" : stderr
            throw MacVMError.appInstallFailed("installing packages into \"\(label)\": \(detail)")
        }
        return MacVMPackage.installedIDs(fromOutput: stdoutText)
    }
    #endif
}
