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

/// **Copy user-provided `.app` bundles into a guest over vsock** — the runtime counterpart to baking
/// them into the base at build time (`scripts/provision-macos-guest.sh` §7½). Two entry points share
/// one streaming core:
///
///  * ``installApps(name:appPaths:)`` targets a RUNNING session clone (Settings → "Copy to running
///    VMs"), so a live VM gets the apps without a restart.
///  * ``installAppsIntoBase(appPaths:)`` briefly boots the **golden base** writable, copies the apps
///    into it, and powers it back off — so newly-added apps persist into every FUTURE clone WITHOUT the
///    ~30-minute reinstall+reprovision a full rebuild costs (Settings → "Add to base image").
///
/// The transfer itself is network-free: the host `tar`s the bundles and streams the archive straight
/// into a guest `tar -x` over the native agent's `exec` stdin (``MacVMExecChannel/writeStdinRaw``).
extension MacVMEngine {
    /// The `.app` bundle paths in `paths` that still exist on disk as directories (a moved/deleted
    /// bundle, or a non-`.app` path, is silently dropped) — the cheap pre-flight both entry points use.
    static func existingAppBundles(_ paths: [String]) -> [String] {
        let fm = FileManager.default
        return paths.filter { path in
            var isDir: ObjCBool = false
            return path.hasSuffix(".app")
                && fm.fileExists(atPath: path, isDirectory: &isDir) && isDir.boolValue
        }
    }

    /// Concatenate `a` then `b`, dropping exact-string duplicates while preserving first-seen order.
    /// Used to union a live Settings selection (apps/packages) with the set carried forward from a base
    /// being replaced, so the settings order wins and a value present in both appears once.
    static func orderedUnion(_ a: [String], _ b: [String]) -> [String] {
        var seen = Set<String>()
        var out: [String] = []
        for value in a + b where seen.insert(value).inserted { out.append(value) }
        return out
    }

    /// Keep the first path for each distinct `.app` **leaf name** (e.g. `Foo.app`), dropping later
    /// paths that would land on the same `/Applications/<leaf>`. Both the base staging (build time) and
    /// the streaming injector key the destination on the leaf, so two source paths sharing a leaf would
    /// otherwise collide; de-duping here keeps the first (Settings-order) one.
    static func dedupedByLeaf(_ paths: [String]) -> [String] {
        var seen = Set<String>()
        var out: [String] = []
        for path in paths where seen.insert(URL(fileURLWithPath: path).lastPathComponent).inserted {
            out.append(path)
        }
        return out
    }

    /// Where a FULL Xcode lives on the host, most-preferred first. Stable before beta: both build a
    /// macOS-27 guest fine, and the stable one is the less surprising thing to bake in.
    static let hostXcodeSearchPaths = ["/Applications/Xcode.app", "/Applications/Xcode-beta.app"]

    /// The first path in `candidates` that is a real full Xcode — i.e. carries an `xcodebuild`, which
    /// the Command Line Tools do **not** ship at all. That binary is the whole point: it's the only
    /// supported installer for the Metal toolchain (`xcodebuild -downloadComponent MetalToolchain`),
    /// and the Metal compiler only resolves through `xcrun` when a full Xcode is the selected developer
    /// dir — so a CLT-only guest can neither install nor use `metal`.
    static func preferredHostXcode(
        _ candidates: [String] = hostXcodeSearchPaths,
        exists: (String) -> Bool = { FileManager.default.fileExists(atPath: $0) }
    ) -> String? {
        candidates.first { exists("\($0)/Contents/Developer/usr/bin/xcodebuild") }
    }

    /// Append the host's full Xcode to the `.app` bundles being baked into the base, so every session
    /// clone is born able to compile Metal (see ``preferredHostXcode(_:exists:)`` and the provisioner's
    /// Phase 7⅞). No host Xcode ⇒ unchanged, and an operator who already staged an `Xcode*.app` of
    /// their own wins — we never stage a second Xcode alongside theirs.
    static func addingHostXcode(_ staged: [String], hostXcode: String?) -> [String] {
        guard let hostXcode else { return staged }
        let alreadyStaged = staged.contains { path in
            let leaf = URL(fileURLWithPath: path).lastPathComponent
            return leaf.hasPrefix("Xcode") && leaf.hasSuffix(".app")
        }
        return alreadyStaged ? staged : staged + [hostXcode]
    }

    /// Fold app host paths just injected into `bundle` at runtime (``installAppsIntoBase``) into its
    /// recorded ``MacVMBaseStatus/installedAppPaths``, so a later from-scratch rebuild that replaces
    /// this base carries them forward. Union preserves the existing order; new paths append.
    func recordBaseInstalledApps(_ bundle: MacVMBundle, appPaths: [String]) {
        guard !appPaths.isEmpty else { return }
        var status = readBaseStatus(bundle)
        status.installedAppPaths = Self.orderedUnion(status.installedAppPaths, appPaths)
        writeBaseStatus(bundle, status)
    }

    /// The package-injection twin of ``recordBaseInstalledApps(_:appPaths:)`` — folds package ids just
    /// installed into `bundle` (``installPackagesIntoBase``) into ``MacVMBaseStatus/installedPackageIDs``.
    func recordBaseInstalledPackages(_ bundle: MacVMBundle, packageIDs: [String]) {
        guard !packageIDs.isEmpty else { return }
        var status = readBaseStatus(bundle)
        status.installedPackageIDs = Self.orderedUnion(status.installedPackageIDs, packageIDs)
        writeBaseStatus(bundle, status)
    }

    /// Copy each `.app` in `appPaths` into the named **running** guest's `/Applications`, replacing any
    /// existing copy and stripping quarantine. Returns the bundle names that landed. Throws
    /// ``MacVMError/notRunning(_:)`` if the VM isn't live, or ``MacVMError/appInstallFailed(_:)`` on a
    /// transfer failure. An empty/all-invalid set is a no-op returning `[]`.
    @discardableResult
    public func installApps(name: String, appPaths: [String]) async throws -> [String] {
        #if arch(arm64)
        guard let instance = live[name]?.instance else { throw MacVMError.notRunning(name) }
        return try await installApps(onInstance: instance, label: name, appPaths: appPaths)
        #else
        throw MacVMError.unavailable(Self.unsupportedReason ?? "requires Apple silicon")
        #endif
    }

    /// Copy the given `.app` bundles 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), stream the apps
    /// into `/Applications`, then shut it down cleanly so the quiesced disk can be cloned again. New
    /// session VMs cut from the base afterward include the apps. Returns the bundle names that landed.
    ///
    /// Refuses (``MacVMError/baseBusy(_:)``) while the base is otherwise in use (a build, a Recovery
    /// session, or a concurrent app injection), 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 by running sessions).
    @discardableResult
    public func installAppsIntoBase(appPaths: [String]) async throws -> [String] {
        #if arch(arm64)
        guard !baseIsBusy else {
            throw MacVMError.baseBusy(
                baseBuilding
                    ? "a base image build is in progress"
                    : "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 apps to it.")
        }
        let apps = Self.existingAppBundles(appPaths)
        guard !apps.isEmpty else { return [] }

        // 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 apps to the running VMs instead.")
        }

        baseAppInstalling = true
        baseWorkCancelled = false  // a new, deliberate pass — re-enable bundle writes
        // Hold the anti-nap assertion for the boot (the base VM never enters `live`, which is what the
        // assertion otherwise keys off) so a backgrounded app isn't App-Napped mid-injection.
        syncBackgroundActivity()
        // Publish the base as a running, system-managed VM so the Control panel's Virtual Machines
        // section shows it (with a live activity note) alongside the per-session guests while apps
        // install — the user's indication that the base is booted and working.
        baseMaintenance = MacVMMaintenanceInfo(
            name: "base-app-install", os: .macOS, activity: "Starting…", booting: true)
        defer {
            baseAppInstalling = false
            baseMaintenance = nil
            syncBackgroundActivity()
        }

        let config = try Self.makeConfiguration(
            bundle: base, cpus: MacVMSettings.vmCPUs, memoryGiB: MacVMSettings.vmMemoryGiB,
            mac: Self.randomMAC(), mounts: [])
        let instance = MacVMInstance(configuration: config, label: "base-app-install")
        do {
            try await instance.start()
        } catch {
            throw MacVMError.appInstallFailed("could not boot the base image: \(error)")
        }
        do {
            try await awaitAgentReady(onInstance: instance, label: "base-app-install")
            baseMaintenance = MacVMMaintenanceInfo(
                name: "base-app-install", os: .macOS, activity: "Installing apps…", booting: false)
            // Bound the vsock copy so a wedged guest can't hang the whole operation — and leave the
            // base VM booted — forever. The copy's own stdin pump is already capped at 30 min; this
            // outer deadline also covers the wait for the guest's `exit` frame, so a guest that stops
            // responding after the pump surfaces as a clean, recoverable failure (→ force power-off
            // below) instead of an indefinite stall.
            let installed: [String]
            do {
                installed = try await withAgentTimeout(seconds: 2000) { [apps] in
                    try await self.installApps(onInstance: instance, label: "base", appPaths: apps)
                }
            } catch is MacVMAgentError {  // only the timeout injects this; installApps throws MacVMError
                throw MacVMError.appInstallFailed(
                    "the base image stopped responding while copying apps (timed out). It has been "
                        + "powered off — try again.")
            }
            await shutdownBaseGuest(instance)  // clean power-off so the disk quiesces before re-cloning
            // Remember these apps live in the base now, so a later rebuild that replaces it carries
            // them forward even if they're no longer in the "Included apps" list.
            recordBaseInstalledApps(base, appPaths: apps)
            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 streaming core: `tar` the bundles on the host and pipe the archive into a guest `tar -x`
    /// over `exec` stdin. `label` names the target only for error text.
    private func installApps(
        onInstance instance: MacVMInstance, label: String, appPaths: [String]
    ) async throws -> [String] {
        // Keep existing `.app` directories, de-duplicated by destination leaf name (last path wins).
        var sourceByLeaf: [String: URL] = [:]
        var order: [String] = []
        for path in Self.existingAppBundles(appPaths) {
            let url = URL(fileURLWithPath: path)
            let leaf = url.lastPathComponent
            if sourceByLeaf[leaf] == nil { order.append(leaf) }
            sourceByLeaf[leaf] = url
        }
        guard !order.isEmpty else { return [] }

        // Host side: `tar -c -z` the bundles to stdout, each rooted at the archive top via a per-app
        // `-C <parent> <leaf>` (bsdtar applies `-C` positionally to the file that follows).
        let tar = Process()
        tar.executableURL = URL(fileURLWithPath: "/usr/bin/tar")
        var args = ["-c", "-z", "-f", "-"]
        for leaf in order {
            let src = sourceByLeaf[leaf]!
            args.append(contentsOf: ["-C", src.deletingLastPathComponent().path, leaf])
        }
        tar.arguments = args
        let tarOut = Pipe()
        tar.standardOutput = tarOut
        tar.standardError = Pipe()
        do {
            try tar.run()
        } catch {
            throw MacVMError.appInstallFailed("could not read the app bundles on the host: \(error)")
        }

        // Guest side: drop any existing copy, unpack the streamed archive into /Applications preserving
        // permissions, then strip quarantine. The agent account is an admin and /Applications is
        // admin-writable, so no sudo (there's no tty/password on a live guest) is needed.
        let quotedNames = order.map(Self.shQuote).joined(separator: " ")
        let remoteBody = """
        set -e
        DEST=/Applications
        for n in \(quotedNames); do /bin/rm -rf "$DEST/$n"; done
        /usr/bin/tar -x -z -p -f - -C "$DEST"
        for n in \(quotedNames); do /usr/bin/xattr -dr com.apple.quarantine "$DEST/$n" 2>/dev/null || true; done
        """

        let channel: MacVMExecChannel
        do {
            channel = try await openExecChannel(
                instance: instance, label: label, workdir: nil, env: [:], remoteBody: remoteBody)
        } catch {
            tar.terminate()
            throw MacVMError.appInstallFailed("\(error)")
        }

        // Pump the tar archive into the guest's stdin on a detached task: `writeStdinRaw` blocks on
        // guest backpressure (a slow unpack), which would otherwise park the actor. Natural end-to-end
        // flow control — a full pipe stalls the host `tar` too — so no unbounded buffering.
        let reader = tarOut.fileHandleForReading
        let streamTask = Task.detached(priority: .utility) { () -> Bool in
            // ALWAYS signal stdin EOF, even on an early exit (deadline/backpressure/torn channel).
            // Without it the guest `tar -x` blocks forever waiting for more input, so the guest shell
            // never exits and never emits its `exit` frame — which would hang `wait()`/the stdout
            // drain below (and leave the base VM booted) indefinitely. A truncated archive instead
            // fails `tar` cleanly, surfacing as a nonzero exit rather than a permanent stall.
            defer { channel.closeStdin() }
            let deadline = Date().addingTimeInterval(1800)  // ≤30 min even for a very large bundle set
            while true {
                let chunk = reader.availableData
                if chunk.isEmpty { break }  // host tar closed stdout — archive complete
                if !channel.writeStdinRaw(chunk, deadline: deadline) { return false }
            }
            return true
        }

        // Drain both guest streams concurrently so the exit frame is processed and stderr is captured.
        async let stderrText = Self.drainLines(channel.stderrLines)
        _ = await Self.drainLines(channel.stdoutLines)
        let streamed = await streamTask.value
        let code = await channel.wait()
        tar.waitUntilExit()
        let stderr = await stderrText

        guard streamed, code == 0, tar.terminationStatus == 0 else {
            let detail = stderr.isEmpty ? "guest exit \(code)" : stderr
            throw MacVMError.appInstallFailed("into \"\(label)\": \(detail)")
        }
        return order
    }

    /// Wait (bounded) until the agent in `instance` answers a vsock `ping` — the instance-based twin of
    /// ``awaitAgentReady(name:)`` for a VM booted outside the `live` registry (the base injection).
    /// Shared by the app and common-package base injectors (`MacVMEngine+InstallPackages.swift`).
    ///
    /// `label` names the pass that is waiting. The failure is `agentUnavailable`, not
    /// `appInstallFailed`: every caller but the app injector is doing something else entirely, and an
    /// MDM pass reporting "the base image's in-guest agent never answered" as an *app install* failure
    /// sent the last investigation looking in the wrong file. The overwhelmingly common cause is that
    /// the guest did not auto-login — the agent's LaunchAgent is `LimitLoadToSessionType: Aqua`, so it
    /// does not exist at the login window — and the message now says so.
    func awaitAgentReady(
        onInstance instance: MacVMInstance, label: String = "base",
        timeout: TimeInterval = 300
    ) async throws {
        let deadline = Date().addingTimeInterval(timeout)
        while Date() < deadline {
            if let box = try? await withAgentTimeout(seconds: 5, {
                try await instance.connectAgent(port: MacVMAgentWire.port)
            }) {
                let client = MacVMAgentClient(connection: box)
                if let pong = try? await client.request(op: "ping", timeout: 5),
                    pong["version"]?.intValue == MacVMAgentWire.version
                {
                    await client.close()
                    return
                }
                await client.close()
            }
            try? await Task.sleep(nanoseconds: 2_000_000_000)
        }
        throw MacVMError.agentUnavailable(
            "the in-guest agent never answered a vsock ping during \"\(label)\" within "
                + "\(Int(timeout))s. The agent runs as an Aqua LaunchAgent, so the usual cause is a "
                + "guest that did not auto-login into the agent account.")
    }

    /// Cleanly power the injection guest off (so the base disk quiesces before it's cloned again):
    /// ask the guest to `shutdown -h` via the agent account's password, then wait for it to stop,
    /// force-stopping only if it overruns. Best effort — a torn disk is the thing we're avoiding.
    /// Shared by the app and common-package base injectors (`MacVMEngine+InstallPackages.swift`).
    ///
    /// This is the last shutdown before the base is cloned, so it also **quits Terminal and clears its
    /// saved window state** first. The base is built by HID-driving Terminal (`typeBootstrapLaunch`),
    /// and macOS's window restoration otherwise saves that Terminal window into the sealed image and
    /// reopens it on every clone boot — cluttering the desktop we screenshot. Killing Terminal (no
    /// AppleEvents/TCC needed) and removing its `savedState`, with restoration pinned off, seals a
    /// clean desktop. Runs as the agent user over vsock, so it never touches this exec's own process.
    func shutdownBaseGuest(_ instance: MacVMInstance) async {
        if let password = try? resolveAgentPassword() {
            let cleanup =
                "/usr/bin/defaults write com.apple.Terminal NSQuitAlwaysKeepsWindows -bool false 2>/dev/null || true; "
                + "/usr/bin/defaults write com.apple.loginwindow TALLogoutSavesState -bool false 2>/dev/null || true; "
                + "/usr/bin/killall Terminal 2>/dev/null || true; /bin/sleep 1; "
                + "/bin/rm -rf \"$HOME/Library/Saved Application State/com.apple.Terminal.savedState\" 2>/dev/null || true; "
            let body = cleanup
                + "printf '%s\\n' \(Self.shQuote(password)) | /usr/bin/sudo -S -k /sbin/shutdown -h now"
            // Fire-and-forget: the connection dies as the guest powers down, so don't await an exit.
            if let channel = try? await openExecChannel(
                instance: instance, label: "base", workdir: nil, env: [:], remoteBody: body)
            {
                channel.closeStdin()
            }
        }
        await instance.awaitStopped(timeout: 120)
    }

    /// Drain an exec line stream into one newline-joined string, capped so a runaway guest can't
    /// balloon host memory (our command emits little, but the cap keeps it honest).
    static func drainLines(_ stream: AsyncThrowingStream<Data, Error>, cap: Int = 64 * 1024) async
        -> String
    {
        var lines: [String] = []
        var bytes = 0
        do {
            for try await line in stream where bytes < cap {
                lines.append(String(decoding: line, as: UTF8.self))
                bytes += line.count + 1
            }
        } catch { /* stream ended on error — return what we have */ }
        return lines.joined(separator: "\n")
    }
    #endif
}
