nucleic/Sources/NucleicCore/Container/ContainerizedProcessHandle.swift

import Foundation
import Containerization

/// Bridges a containerization-framework `LinuxProcess` (a process running inside the guest VM,
/// with stdio carried over vsock) onto Nucleic's existing ``ProcessHandle`` contract — the same
/// protocol a host-spawned ``ChildProcess`` satisfies. Because both conform to ``ProcessHandle``,
/// every backend (``ClaudeCodeBackend`` and friends) drives an in-container agent with the *exact*
/// same code path it uses for a host spawn: `writeLine` the NDJSON prompt, iterate `stdoutLines`,
/// `wait()` for the exit code, `sendSignal` to interrupt.
///
/// The framework's stdio model is: you attach a ``Writer`` for stdout/stderr and a ``ReaderStream``
/// for stdin on the `LinuxProcessConfiguration` *before* the process starts (there is no
/// re-attaching afterward). ``StreamWriter``/``StreamReader`` below are those attachments; this
/// handle wires them to async line streams and a stdin pump that frame bytes identically to the
/// host path (one newline-delimited line at a time, via the shared ``LineSplitter``).
public final class ContainerizedProcessHandle: ProcessHandle, @unchecked Sendable {
    public let stdoutLines: AsyncThrowingStream<Data, Error>
    public let stderrLines: AsyncThrowingStream<Data, Error>

    private let process: LinuxProcess
    private let stdin: StreamReader
    private let exitTask: Task<Int32, Never>

    /// The guest process id. Best-effort: the framework exposes it once started.
    public var processID: Int32 { Int32(process.pid) }

    /// Built by ``ContainerEngine/exec(name:workdir:env:argv:uid:gid:)`` *after* `process.start()`,
    /// handing in the same stdin reader, stdout/stderr line streams, and the ``StreamWriter``s that
    /// were attached to the process configuration. Spins up the exit waiter immediately so `wait()`
    /// is cheap to await repeatedly.
    init(
        process: LinuxProcess,
        stdin: StreamReader,
        stdoutWriter: StreamWriter,
        stderrWriter: StreamWriter,
        stdoutLines: AsyncThrowingStream<Data, Error>,
        stderrLines: AsyncThrowingStream<Data, Error>
    ) {
        self.process = process
        self.stdin = stdin
        self.stdoutLines = stdoutLines
        self.stderrLines = stderrLines
        // `LinuxProcess.wait()` resolves once the guest process exits; cache the code so the backend
        // can await it repeatedly (decode loop, then shutdown). Crucially, the framework does NOT
        // call `Writer.close()` at guest-stdout EOF — but `wait()` returns only AFTER stdio fully
        // drains (its `waitIoComplete`), so closing the writers here flushes any trailing partial
        // line and FINISHES the line streams. Without this the backend's `for await stdoutLines`
        // loop would never end (it relies on EOF, exactly like the host `ChildProcess` path).
        self.exitTask = Task {
            let status = try? await process.wait()
            try? stdoutWriter.close()
            try? stderrWriter.close()
            return status?.exitCode ?? -1
        }
    }

    /// Frame and forward one NDJSON line to the guest's stdin (newline appended to match
    /// ``ChildProcess/writeLine(_:)``). The agent reads `--input-format stream-json` off stdin.
    public func writeLine(_ data: Data) throws {
        stdin.send(data)
    }

    /// Signal guest-stdin EOF — the agent's single-shot run ends its input phase. Mirrors closing
    /// the host stdin pipe.
    public func closeStdin() {
        stdin.finish()
    }

    /// Deliver a POSIX signal to the guest process (SIGINT for Stop, SIGTERM on shutdown). The
    /// framework's `kill` is async/throws; fire-and-forget from this synchronous contract method.
    public func sendSignal(_ sig: Int32) {
        let process = self.process
        Task { try? await process.kill(Signal(rawValue: sig)) }
    }

    /// The guest process's exit code (or -1 if the wait failed). Drives the backend's
    /// exit-handling / OOM-diagnosis path exactly as a host exit status does.
    public func wait() async -> Int32 {
        await exitTask.value
    }
}

/// stdout/stderr sink handed to a `LinuxProcessConfiguration`. The framework calls `write(_:)` with
/// each raw byte chunk the guest process emits; we run it through the shared ``LineSplitter`` so the
/// backend receives one newline-stripped line per stream element — byte-for-byte the same framing as
/// the host ``ChildProcess`` line stream. `close()` (guest stream EOF) drains the trailing partial
/// line and finishes the async stream.
final class StreamWriter: Containerization.Writer, @unchecked Sendable {
    private let splitter = LineSplitter()
    private let continuation: AsyncThrowingStream<Data, Error>.Continuation

    init(_ continuation: AsyncThrowingStream<Data, Error>.Continuation) {
        self.continuation = continuation
    }

    func write(_ data: Data) throws {
        for line in splitter.push(data) { continuation.yield(line) }
    }

    func close() throws {
        if let rest = splitter.drain() { continuation.yield(rest) }
        continuation.finish()
    }
}

/// stdin source handed to a `LinuxProcessConfiguration`. The framework iterates `stream()` and
/// forwards each `Data` chunk to the guest process's stdin. ``send(_:)`` appends a newline so each
/// call writes exactly one NDJSON line (matching ``ChildProcess/writeLine(_:)``); ``finish()`` ends
/// the stream, which the framework translates into guest-stdin EOF.
final class StreamReader: Containerization.ReaderStream, @unchecked Sendable {
    private let backing: AsyncStream<Data>
    private let continuation: AsyncStream<Data>.Continuation

    init() {
        (backing, continuation) = AsyncStream<Data>.makeStream()
    }

    func stream() -> AsyncStream<Data> { backing }

    func send(_ data: Data) {
        var framed = data
        framed.append(0x0A)
        continuation.yield(framed)
    }

    func finish() {
        continuation.finish()
    }
}