nucleic/Sources/NucleicCore/Container/ContainerManager.swift

import Foundation

/// App-level coordinator for per-session sandbox containers (injected like
/// `GitWorktreeManager`). Owns the lifecycle the `ContainerEngine` is too low-level to
/// track: keeping one long-lived container per session, an idle timer that stops it after
/// inactivity, teardown on session end, and on-disk reconciliation on launch.
///
/// Activity is reference-counted: `ensureRunning` (called at the start of every turn)
/// marks the container busy and disarms the idle timer; `finished` marks the turn done and
/// re-arms it. The container is only stopped when no turn is in flight, so a single
/// long-running turn is never killed mid-flight.
public actor ContainerManager {
    private let engine: ContainerEngine

    /// Live per-container bookkeeping, keyed by **logical** container name (the stable name in this
    /// type's public API — `containerName(for:)` / `sharedContainerName(for:split:)`).
    private var active: [String: Int] = [:]              // in-flight turn count
    private var idleTimeouts: [String: TimeInterval] = [:]
    private var idleTimers: [String: Task<Void, Never>] = [:]

    /// Logical → **physical** container name for this run. The shared per-family control containers
    /// run under a randomized physical name (a stable base + a random fun word, re-rolled each time
    /// the sandbox is created) so an agent that glimpses another sandbox — via `ps`, a stray pid, a
    /// host-exec — can neither recognize it as a sibling nor guess/target it by name. Bookkeeping and
    /// every public method speak logical names; this map is applied only at the engine boundary.
    /// Non-randomized names (the primary shared container, per-session containers) map to themselves.
    /// Populated lazily on `ensureRunning`; an entry is dropped on teardown/recreate so the next
    /// creation re-randomizes, and the whole map resets each app launch.
    private var physicalNames: [String: String] = [:]

    /// Invoked with a shared control container's **logical** name the moment that container is
    /// *permanently* removed (``teardownShared()`` / ``recreateShared()``), so the owner can release
    /// the container-scoped approval server and its control socket. Idle-stop / restart deliberately do
    /// NOT fire it: the long-lived server is meant to survive those, and the reconnecting session's
    /// `MCPApprovalServer.start(unixSocketPath:)` re-validates the socket, so tearing it down there
    /// would be pure churn. Nil in tests / host-only setups.
    private let onSharedContainerRemoved: (@Sendable (String) async -> Void)?

    public init(
        engine: ContainerEngine = ContainerEngine(),
        onSharedContainerRemoved: (@Sendable (String) async -> Void)? = nil
    ) {
        self.engine = engine
        self.onSharedContainerRemoved = onSharedContainerRemoved
    }

    /// Per-channel suffix appended to every container name on non-release builds, so several
    /// installed Nucleic builds (release + beta + local dev, etc.) running against the **shared**
    /// on-disk container store (`ContainerEngine.defaultStorageRoot`) don't collide on their
    /// per-container rootfs clones. Empty for the stable release; `-canary` / `-beta` / `-rc` /
    /// `-local` for the non-release channels. Baked in at build time from the `NUCLEIC_CHANNEL` define
    /// (mirrors NucleicApp's `BuildChannel`); NucleicCore is compiled with that define — see Package.swift.
    public nonisolated static let channelSuffix: String = {
        #if NUCLEIC_STABLE
        return ""
        #elseif NUCLEIC_RC
        return "-rc"
        #elseif NUCLEIC_BETA
        return "-beta"
        #elseif NUCLEIC_CANARY
        return "-canary"
        #else
        return "-local"
        #endif
    }()

    /// Every non-empty `channelSuffix` a build can carry. Used to tell a stable build's untagged
    /// containers apart from other channels' tagged ones during disk GC. Keep in lockstep with the
    /// `channelSuffix` cases above.
    public nonisolated static let nonReleaseSuffixes = ["-local", "-canary", "-beta", "-rc"]

    /// Whether a container/clone name belongs to THIS build channel's namespace. The on-disk
    /// container store is shared across installed builds (`ContainerEngine.defaultStorageRoot`), so
    /// the launch-time disk GC uses this to reap only its own channel's clones and never another
    /// build's. A non-release build owns names ending in its `channelSuffix`; the stable release owns
    /// names carrying none of the non-release tags.
    public nonisolated static func ownsContainer(named name: String) -> Bool {
        channelSuffix.isEmpty
            ? !nonReleaseSuffixes.contains(where: name.hasSuffix)
            : name.hasSuffix(channelSuffix)
    }

    /// Stable container name for a session. Lowercased because `UUID().uuidString` (hence
    /// `SessionID.short`) is uppercase, while container names are conventionally lowercase
    /// and our reconcile parser matches lowercase hex. Carries the per-channel `channelSuffix`.
    public nonisolated static func containerName(for session: SessionID) -> String {
        "nucleic-\(session.short.lowercased())\(channelSuffix)"
    }

    /// The single **primary** container shared by all Nucleic Control projects that haven't
    /// opted into per-session containers. Deliberately not of the `nucleic-<8 hex>` shape, so
    /// the reconcile/list parser never mistakes it for a per-session container (and so a
    /// per-session `teardown` — which computes a `nucleic-<short>` name — never removes it).
    /// Its busy/idle lifecycle is ref-counted by name across every control session that uses it.
    /// Used by every backend when backend-splitting is off; with splitting on, each agent family
    /// gets its own suffixed sibling (`-claude`/`-codex`/`-xai`) instead. Carries the per-channel
    /// `channelSuffix`.
    public static let sharedControlContainerName = "nucleic-control" + channelSuffix

    /// The per-family shared control containers, used only when
    /// `ContainerServiceSettings.splitControlContainersByBackend` is on — one sandbox per agent
    /// family so competitive Claude / GPT / xAI agents can't kill each other's processes
    /// ("agenticide") in a shared box. Same non-`nucleic-<8 hex>` shape as the primary, so they're
    /// likewise invisible to the per-session list parser. Each carries the `channelSuffix` at the
    /// very end (e.g. `nucleic-control-codex-beta`).
    public static let claudeControlContainerName = "nucleic-control-claude" + channelSuffix
    public static let codexControlContainerName = "nucleic-control-codex" + channelSuffix
    public static let xaiControlContainerName = "nucleic-control-xai" + channelSuffix
    public static let opencodeControlContainerName = "nucleic-control-opencode" + channelSuffix
    public static let hermesControlContainerName = "nucleic-control-hermes" + channelSuffix
    public static let cursorControlContainerName = "nucleic-control-cursor" + channelSuffix

    /// Every shared control container name that can exist. Lifecycle sweeps (reconcile, teardown,
    /// status, recreate) iterate this so they stay correct no matter how the split setting is
    /// currently configured — e.g. turning splitting back off still reaps the per-family containers.
    public static let allSharedControlContainerNames =
        [sharedControlContainerName, claudeControlContainerName,
         codexControlContainerName, xaiControlContainerName,
         opencodeControlContainerName, hermesControlContainerName, cursorControlContainerName]

    /// The shared control container a session should use, given its backend and whether
    /// backend-splitting is on. With splitting off, every backend shares `nucleic-control`; with it
    /// on, each agent family gets its own sandbox: Claude → `nucleic-control-claude`, GPT/Codex →
    /// `nucleic-control-codex`, xAI/Grok → `nucleic-control-xai`.
    public nonisolated static func sharedContainerName(
        for backend: BackendID, split: Bool
    ) -> String {
        guard split else { return sharedControlContainerName }
        switch backend {
        case .claudeCode: return claudeControlContainerName
        case .codex, .codexExec: return codexControlContainerName
        case .grok: return xaiControlContainerName
        case .opencode: return opencodeControlContainerName
        case .hermes: return hermesControlContainerName
        case .cursorAgent: return cursorControlContainerName
        }
    }

    // MARK: - Randomized physical names

    /// The logical control-container names that run under a **randomized** physical name: the
    /// per-agent-family containers. The primary `sharedControlContainerName` is excluded — when
    /// splitting is off it's the lone control sandbox, so there's no sibling for an agent to target.
    public nonisolated static let randomizedControlNames: Set<String> =
        [claudeControlContainerName, codexControlContainerName, xaiControlContainerName,
         opencodeControlContainerName, hermesControlContainerName, cursorControlContainerName]

    /// Friendly type label for a shared control container's *logical* name (Control panel heading).
    public nonisolated static func controlTypeLabel(forLogical logical: String) -> String {
        switch logical {
        case claudeControlContainerName: return "Claude"
        case codexControlContainerName: return "Codex"
        case xaiControlContainerName: return "xAI"
        case opencodeControlContainerName: return "OpenCode"
        case hermesControlContainerName: return "Hermes"
        case cursorControlContainerName: return "Cursor"
        default: return "Shared"
        }
    }

    /// Dictionary of fun / curious / odd words the random sandbox suffix is drawn from. Lowercase,
    /// single tokens (container names are lowercase, hyphen-delimited), and deliberately free of the
    /// channel tags (`local`/`canary`/`beta`/`rc`) so a suffix can never look like a `channelSuffix` and
    /// confuse the disk-GC's channel scoping.
    public nonisolated static let funWords: [String] = [
        "wobble", "kerfuffle", "noodle", "quokka", "zephyr", "pickle", "bumble", "gizmo",
        "waffle", "snazzy", "doodad", "kumquat", "flapjack", "wombat", "gadget", "muffin",
        "pretzel", "narwhal", "tangelo", "bramble", "fiddle", "galoot", "hodgepodge", "jamboree",
        "kazoo", "lollop", "moxie", "nincompoop", "oodles", "persnickety", "quibble", "razzle",
        "skedaddle", "tumbleweed", "umpteen", "vortex", "whirligig", "yonder", "zigzag", "bonbon",
        "cattywampus", "discombobulate", "doozy", "flummox", "gobbledygook", "hullabaloo", "lollygag",
        "malarkey", "nifty", "obelisk", "pumpernickel", "quagmire", "rumpus", "snickerdoodle",
        "thingamajig", "widget", "zonk", "blunderbuss", "cahoots", "dingus", "epiphany", "frippery",
        "gubbins", "haberdash", "iguana", "junket", "kerplunk", "limerick", "mollusk", "nimbus",
        "ottoman", "platypus", "quasar", "rutabaga", "squiggle", "trundle", "ukulele", "vagabond",
        "walrus", "xylophone", "yodel", "ziggurat", "abacus", "bazooka", "conundrum", "dapper",
        "embiggen", "flabbergast", "gizzard", "hobnob", "inkling", "jubilee", "kerchief", "loofah",
        "marmalade", "nugget", "octopus", "paprika", "quokkas", "ricochet", "sasquatch", "tadpole",
        "umbrella", "verbena", "whatsit", "yowza", "zucchini", "bumblebee", "cucumber", "doohickey",
        "fandango", "gargoyle", "hiccup", "jellybean", "kookaburra", "lemur", "meerkat", "noggin",
    ]

    /// Pick a random fun word for a new sandbox.
    public nonisolated static func randomFunWord() -> String { funWords.randomElement() ?? "sandbox" }

    /// Compose the physical name for a randomized control container from a fun `word`:
    /// `nucleic-control-<word><channelSuffix>`. Crucially the agent *family is not encoded* — all
    /// split families share the `nucleic-control-…` shape — so an agent that glimpses a sibling's
    /// name can't tell which family (Claude / Codex / xAI) it serves, only the trusted host UI maps
    /// name → type. The channel tag stays last so the disk-GC's `ownsContainer` scoping still works.
    /// Pure, for testability.
    public nonisolated static func physicalControlName(word: String) -> String {
        "nucleic-control-\(word)\(channelSuffix)"
    }

    /// The physical name a logical name currently maps to (identity until materialized).
    private func physical(_ logical: String) -> String { physicalNames[logical] ?? logical }

    /// Resolve (allocating on first use) the physical name for `logical`. Randomized control names
    /// get a fresh fun-word physical name — distinct from every other live sandbox's, so the three
    /// split families never collide on a word; everything else maps to itself.
    private func materializePhysical(_ logical: String) -> String {
        if let existing = physicalNames[logical] { return existing }
        guard Self.randomizedControlNames.contains(logical) else {
            physicalNames[logical] = logical
            return logical
        }
        let used = Set(physicalNames.values)
        var phys = Self.physicalControlName(word: Self.randomFunWord())
        var tries = 0
        while used.contains(phys) && tries < 50 {
            phys = Self.physicalControlName(word: Self.randomFunWord())
            tries += 1
        }
        physicalNames[logical] = phys
        return phys
    }

    /// Ensure the session's container is up and return its name plus the host-gateway IP
    /// (the container's view of the Mac, for routing the approval MCP server). Marks a turn
    /// as in-flight; pair with `finished(name:)`.
    public func ensureRunning(_ spec: ContainerSpec) async throws
        -> (name: String, hostGateway: String)
    {
        cancelIdleTimer(spec.name)
        active[spec.name, default: 0] += 1
        idleTimeouts[spec.name] = spec.idleTimeout
        do {
            // The engine creates/starts the VM (or returns the already-live one) and hands back
            // the host-gateway address the guest reaches us on — known at interface-allocation
            // time, no in-container `ip route` probe needed. The engine runs it under the randomized
            // physical name; callers keep using the stable logical name (returned here).
            let phys = materializePhysical(spec.name)
            let (_, gateway) = try await engine.ensureRunning(spec.renamed(phys))
            return (spec.name, gateway)
        } catch {
            // Roll back the activity bump so a failed start doesn't pin the container busy.
            finishBookkeeping(spec.name)
            throw error
        }
    }

    /// Run `argv` inside the named container, returning a process handle (forwards to the engine;
    /// no lifecycle effect). The handle streams the guest process's stdio over vsock and satisfies
    /// the same `ProcessHandle` contract a host spawn does.
    public func exec(
        name: String, workdir: String, env: [String: String], argv: [String],
        uid: Int? = nil, gid: Int? = nil
    ) async throws -> any ProcessHandle {
        try await engine.exec(
            name: physical(name), workdir: workdir, env: env, argv: argv, uid: uid, gid: gid)
    }

    /// Mark a turn finished. When the last in-flight turn ends, arm the idle timer.
    public func finished(name: String) {
        finishBookkeeping(name)
    }

    /// Stop a throwaway container the agent created via `linux_container` (keeps its clone on disk so
    /// a later `exec` can restart it). Clears its activity bookkeeping so no idle timer lingers.
    public func stopAgentContainer(name: String) async {
        cancelIdleTimer(name)
        active[name] = nil
        idleTimeouts[name] = nil
        await engine.stop(name: physical(name))
    }

    /// Stop and delete a throwaway container the agent created via `linux_container`. Returns true if
    /// it's gone (or never existed). Clears its bookkeeping and physical-name mapping.
    @discardableResult
    public func removeAgentContainer(name: String) async -> Bool {
        cancelIdleTimer(name)
        active[name] = nil
        idleTimeouts[name] = nil
        let removed = await engine.remove(name: physical(name))
        physicalNames[name] = nil
        return removed
    }

    private func finishBookkeeping(_ name: String) {
        let remaining = max(0, (active[name] ?? 0) - 1)
        active[name] = remaining
        if remaining == 0 { armIdleTimer(name) }
    }

    /// Stop and remove the session's container. When `waitForActive` is true, first wait for
    /// any in-flight turn in that container to finish (graceful — used when sandboxing is
    /// turned off for a project); otherwise remove immediately (archive/delete). Returns the
    /// container name if it could NOT be removed, so the caller can warn the user; `nil` on
    /// success (including when no such container existed).
    @discardableResult
    public func teardown(_ session: SessionID, waitForActive: Bool = false) async -> String? {
        let name = Self.containerName(for: session)
        if waitForActive { await waitUntilIdle(name) }
        cancelIdleTimer(name)
        active[name] = nil
        idleTimeouts[name] = nil
        let removed = await engine.remove(name: physical(name))
        physicalNames[name] = nil
        return removed ? nil : name
    }

    /// Stop and remove the shared primary control container(s) and clear their bookkeeping. Skips
    /// any container with a turn still in flight (so a sibling control session isn't yanked out)
    /// and any that doesn't exist. Sweeps every shared name — both the Claude and the split-out
    /// Codex container — so a session of either family is reaped. Per-session `teardown` can't
    /// reach these (it derives a `nucleic-<short>` name), so AppStore calls this once the last
    /// control session is gone.
    public func teardownShared() async {
        for name in Self.allSharedControlContainerNames {
            guard (active[name] ?? 0) == 0 else { continue }
            cancelIdleTimer(name)
            active[name] = nil
            idleTimeouts[name] = nil
            await engine.remove(name: physical(name))
            physicalNames[name] = nil  // next creation re-randomizes
            await onSharedContainerRemoved?(name)  // release the container-scoped approval server + socket
        }
    }

    /// Health of the shared Nucleic Control container(s), for the Control panel. Reports running /
    /// exists if *any* shared control container (Claude or the split-out Codex one) is up. Best-
    /// effort: a running probe (`exec true`) first — the common case — falling back to an existence
    /// check only when none is running.
    public func controlContainerStatus() async -> (exists: Bool, running: Bool) {
        var exists = false
        for name in Self.allSharedControlContainerNames {
            let phys = physical(name)
            if await engine.isRunning(phys) { return (true, true) }
            if await engine.containerExists(phys) { exists = true }
        }
        return (exists, false)
    }

    /// The shared control containers that currently exist this run, each with its friendly type
    /// label and its **actual (randomized) name**, for the Control panel — which shows the type as
    /// the heading and the opaque name as subtext. Only materialized containers appear.
    public func controlContainers() async -> [ControlContainerEntry] {
        var out: [ControlContainerEntry] = []
        for logical in Self.allSharedControlContainerNames {
            let phys = physical(logical)
            let running = await engine.isRunning(phys)
            var exists = running
            if !exists { exists = await engine.containerExists(phys) }
            guard exists else { continue }
            out.append(ControlContainerEntry(
                typeLabel: Self.controlTypeLabel(forLogical: logical), name: phys, running: running))
        }
        return out
    }

    /// Live CPU/memory usage of the running shared control container, for the Control panel's
    /// resource monitor. Samples the first shared control container that's actually up (Claude's
    /// `nucleic-control` or the split-out Codex one) — the resources are shared, so either reading
    /// reflects the load. Best-effort: `nil` when none is running or the probe fails.
    public func controlContainerUsage() async -> ContainerResourceSample? {
        // Sample each shared name directly — the probe is an in-process `statistics()` read, which
        // returns nil against a stopped/absent container, so no separate running check is needed.
        // The first that yields a reading wins, so the common single-container case costs one read.
        for name in Self.allSharedControlContainerNames {
            if let sample = await engine.sampleResourceUsage(name: physical(name)) { return sample }
        }
        return nil
    }

    /// The shared engine's in-flight artifact download (kernel / init / image pull + unpack), for
    /// the Control panel's progress bar and the chat's delay hint. One engine backs every container,
    /// so this is a single, cheap read (no I/O); `nil` whenever everything is cached and nothing is
    /// downloading.
    public func controlDownloadProgress() async -> ContainerDownloadProgress? {
        await engine.currentDownloadProgress()
    }

    /// Post-mortem for an agent that exited 137/SIGKILL in container `name`: was the container
    /// itself taken down, did the kernel OOM-kill, or was it killed while healthy? Forwards to the
    /// runtime; pure lookup, no lifecycle effect. Best-effort.
    public func diagnoseContainerKill(name: String) async -> ContainerKillDiagnosis {
        await engine.diagnoseKill(name: physical(name))
    }

    /// Suspend until no turn is in flight in `name` (or a safety cap elapses, so a wedged
    /// turn can't block cleanup forever). Awaiting releases the actor, letting `finished`
    /// run and decrement the count.
    private func waitUntilIdle(_ name: String) async {
        let pollNanos: UInt64 = 200_000_000  // 0.2s
        let maxPolls = 3_000                  // ~10 min ceiling
        var polls = 0
        while (active[name] ?? 0) > 0 && polls < maxPolls {
            try? await Task.sleep(nanoseconds: pollNanos)
            polls += 1
        }
    }

    /// On app launch, reconcile on-disk container artifacts. Daemonless: no VM survives the app
    /// process, so there are no live orphans to kill — this is on-disk GC. It keeps only the active
    /// sessions' rootfs clones and drops everything else, including the shared control containers'
    /// clones (so they rebuild fresh on next use, picking up the current bind-mount paths) and image
    /// caches superseded by a bundled-image bump.
    public func reconcile(activeSessions: [SessionID]) async {
        let keep = Set(activeSessions.map { Self.containerName(for: $0) })
        await engine.reconcileDisk(keepNames: keep)
    }

    /// Force a fresh shared Nucleic Control sandbox: stop and remove the shared container and
    /// delete the bundled default image, so the next control session rebuilds the image from the
    /// current Dockerfile and recreates the container. Unconditional (ignores the busy ref-count)
    /// — it's a deliberate "rebuild now" from Settings; a control session mid-turn will have its
    /// container yanked and must retry. Returns once teardown + image delete complete.
    public func recreateShared() async {
        for name in Self.allSharedControlContainerNames {
            cancelIdleTimer(name)
            active[name] = nil
            idleTimeouts[name] = nil
            await engine.remove(name: physical(name))
            physicalNames[name] = nil  // force a fresh random name on the next creation
            await onSharedContainerRemoved?(name)  // release the container-scoped approval server + socket
        }
        await engine.removeDefaultRootfs()
    }

    /// Restart the shared control container(s) in place to reclaim the VM memory the guest holds
    /// onto while running (freed pages aren't returned to the host until the VM is torn down). Keeps
    /// the cached rootfs and the per-container instrumentation (the engine reuses the existing rootfs
    /// clone and skips re-seeding), so it's much lighter than `recreateShared`. Skips any container
    /// with a turn in flight, so live work isn't
    /// yanked — the user can retry once it's idle. Re-arms nothing; the next turn's `ensureRunning`
    /// takes over the idle bookkeeping.
    public func restartShared() async {
        for name in Self.allSharedControlContainerNames {
            guard (active[name] ?? 0) == 0 else { continue }
            cancelIdleTimer(name)
            await engine.restart(name: physical(name))  // same name — reuses the clone
        }
    }

    /// Force an immediate memory-reclaim (balloon) pass across all live containers instead of
    /// waiting for the autoballoon loop's next tick. A no-op for any container whose policy is disabled.
    public func reclaimMemoryNow() async {
        await engine.reclaimMemoryNow()
    }

    // MARK: - Idle timer

    private func armIdleTimer(_ name: String) {
        cancelIdleTimer(name)
        let timeout = idleTimeouts[name] ?? TimeInterval(ProjectSandbox.defaultIdleTimeoutSeconds)
        idleTimers[name] = Task { [weak self] in
            try? await Task.sleep(nanoseconds: UInt64(timeout * 1_000_000_000))
            guard !Task.isCancelled else { return }
            await self?.stopIfIdle(name)
        }
    }

    private func cancelIdleTimer(_ name: String) {
        idleTimers[name]?.cancel()
        idleTimers[name] = nil
    }

    private func stopIfIdle(_ name: String) async {
        idleTimers[name] = nil
        guard (active[name] ?? 0) == 0 else { return }  // a turn snuck in; leave it running
        await engine.stop(name: physical(name))  // keep the mapping — idle-restart reuses the clone
    }
}