import Foundation

// MARK: - ProcessHost contract (ADAPTERS §0.1)

public enum StdinMode: Sendable {
    case pipe
    case closed
}

public struct ProcessSpec: Sendable {
    public let executable: String
    public let args: [String]
    public let cwd: String
    public let env: [String: String]
    public let stdinMode: StdinMode

    public init(
        executable: String, args: [String], cwd: String,
        env: [String: String] = [:], stdinMode: StdinMode = .pipe
    ) {
        self.executable = executable
        self.args = args
        self.cwd = cwd
        self.env = env
        self.stdinMode = stdinMode
    }
}

public protocol ProcessHandle: Sendable {
    /// NDJSON-friendly: one line per element, newline stripped, partial trailing
    /// bytes held across pipe reads.
    var stdoutLines: AsyncThrowingStream<Data, Error> { get }
    var stderrLines: AsyncThrowingStream<Data, Error> { get }
    var processID: Int32 { get }
    func writeLine(_ data: Data) throws
    func closeStdin()
    func sendSignal(_ sig: Int32)
    func wait() async -> Int32
    /// Finish the stdout/stderr line streams immediately, independent of the process's exit, so a
    /// consumer blocked on `for await stdoutLines` is released. Only meaningful when the streams can
    /// outlive a killed process (the containerized handle, whose stream closure rides a guest RPC
    /// that can stall); the host pipe handle EOFs reliably on exit, so its default is a no-op.
    func forceCloseStreams()
}

extension ProcessHandle {
    /// Default: nothing to force. The host `ChildProcess` reads OS pipes that EOF the moment the
    /// process dies, so the line streams always finish on their own. ``ContainerizedProcessHandle``
    /// overrides this because its streams close only when a guest wait RPC resolves.
    public func forceCloseStreams() {}
}

/// One-shot claim guard: the first caller to `claim()` wins, the rest get `false`. Lets two racing
/// tasks (exit vs. timeout) resume a single `CheckedContinuation` exactly once.
private final class ResumeOnceFlag: @unchecked Sendable {
    private let lock = NSLock()
    private var claimed = false
    func claim() -> Bool {
        lock.lock(); defer { lock.unlock() }
        if claimed { return false }
        claimed = true
        return true
    }
}

extension ProcessHandle {
    /// Await the process's exit, but give up after `nanos`. Returns `true` iff it exited in time.
    /// On timeout the background wait is left running (harmless — it resolves once the process
    /// finally dies, e.g. after a follow-up SIGKILL), so this never blocks longer than `nanos`.
    /// Used to bound every Stop/shutdown so a wedged agent can't hang the actor on `wait()`.
    func waitForExit(within nanos: UInt64) async -> Bool {
        let once = ResumeOnceFlag()
        return await withCheckedContinuation { (cont: CheckedContinuation<Bool, Never>) in
            Task { _ = await self.wait(); if once.claim() { cont.resume(returning: true) } }
            Task {
                try? await Task.sleep(nanoseconds: nanos)
                if once.claim() { cont.resume(returning: false) }
            }
        }
    }

    /// Stop the process for good and return once it has actually exited: SIGTERM, then — if it
    /// hasn't gone within `graceNanos` — SIGKILL. For a containerized handle the signals reach the
    /// agent's whole process group. Bounding the wait is what lets a stopped/ended session settle
    /// instead of spinning forever on an agent that ignores the polite signal.
    func terminate(graceNanos: UInt64 = 3 * 1_000_000_000) async {
        sendSignal(SIGTERM)
        if await waitForExit(within: graceNanos) { return }
        sendSignal(SIGKILL)
        // SIGKILL is uncatchable, so the process exits and `wait()` resolves; still bound it so a
        // broken wait can never hang teardown.
        if await waitForExit(within: graceNanos) { return }
        // Killed, but `wait()` still hasn't resolved — the guest wait/stdio RPC is wedged. Finish the
        // line streams ourselves so a consumer blocked on `stdoutLines` is released and teardown ends.
        forceCloseStreams()
    }

    /// Best-effort escalation for an interactive Stop: SIGINT now (let the agent end its turn
    /// cleanly / emit an interrupted result), then SIGKILL the process group if it's still alive
    /// after `graceNanos`. Detached so the caller's `interrupt()` stays non-blocking; the run loop
    /// settles when the process finally exits. A responsive agent exits on SIGINT and the escalation
    /// no-ops; a wedged one is force-killed so Stop always takes effect.
    func interruptThenKill(graceNanos: UInt64 = 5 * 1_000_000_000) {
        sendSignal(SIGINT)
        Task { [self] in
            if await waitForExit(within: graceNanos) { return }
            sendSignal(SIGKILL)
            // If even the kill doesn't settle `wait()` (wedged guest RPC), release any consumer
            // blocked on the line streams so Stop always takes effect.
            if await waitForExit(within: graceNanos) { return }
            forceCloseStreams()
        }
    }

    /// Safety net behind a *cooperative* cancel (e.g. an `turn/interrupt` / `session/cancel` RPC the
    /// agent might ignore): force-SIGKILL the process group if it's still alive after `graceNanos`,
    /// without sending an upfront signal that could disrupt a clean cancel in flight. Detached and a
    /// no-op if the process exits on its own — so a well-behaved agent stops gracefully and a wedged
    /// one is still guaranteed to stop.
    func killGroupAfter(graceNanos: UInt64) {
        Task { [self] in
            if await waitForExit(within: graceNanos) { return }
            sendSignal(SIGKILL)
            if await waitForExit(within: graceNanos) { return }
            forceCloseStreams()
        }
    }
}

/// Spawns and monitors child processes with line-buffered stdio.
public actor ProcessHost {
    public init() {}

    public func launch(_ spec: ProcessSpec) throws -> any ProcessHandle {
        try ChildProcess(spec: spec)
    }
}

// MARK: - Implementation

/// All mutable state is confined: Process control calls (signal/terminate) are
/// thread-safe, stdin writes are serialized behind a lock, and the line streams
/// are produced by dedicated reader tasks.
final class ChildProcess: ProcessHandle, @unchecked Sendable {
    private let process: Process
    private let stdinHandle: FileHandle?
    private let stdinLock = NSLock()
    private var stdinClosed = false

    let stdoutLines: AsyncThrowingStream<Data, Error>
    let stderrLines: AsyncThrowingStream<Data, Error>
    private let exitTask: Task<Int32, Never>

    var processID: Int32 { process.processIdentifier }

    init(spec: ProcessSpec) throws {
        let process = Process()
        if spec.executable.contains("/") {
            process.executableURL = URL(fileURLWithPath: spec.executable)
            process.arguments = spec.args
        } else {
            // Resolve via PATH.
            process.executableURL = URL(fileURLWithPath: "/usr/bin/env")
            process.arguments = [spec.executable] + spec.args
        }
        process.currentDirectoryURL = URL(fileURLWithPath: spec.cwd)
        process.environment = ProcessInfo.processInfo.environment.merging(spec.env) { _, new in new }

        let stdoutPipe = Pipe()
        let stderrPipe = Pipe()
        process.standardOutput = stdoutPipe
        process.standardError = stderrPipe

        switch spec.stdinMode {
        case .pipe:
            let stdinPipe = Pipe()
            process.standardInput = stdinPipe
            stdinHandle = stdinPipe.fileHandleForWriting
        case .closed:
            process.standardInput = FileHandle.nullDevice
            stdinHandle = nil
        }

        let exitStream = AsyncStream<Int32> { continuation in
            process.terminationHandler = { p in
                continuation.yield(p.terminationStatus)
                continuation.finish()
            }
        }
        exitTask = Task {
            for await code in exitStream { return code }
            return -1
        }

        stdoutLines = Self.lineStream(from: stdoutPipe.fileHandleForReading)
        stderrLines = Self.lineStream(from: stderrPipe.fileHandleForReading)

        do {
            try process.run()
        } catch {
            throw BackendError.spawnFailed("\(spec.executable): \(error.localizedDescription)")
        }
        self.process = process
    }

    private static func lineStream(from handle: FileHandle) -> AsyncThrowingStream<Data, Error> {
        AsyncThrowingStream { continuation in
            // Chunked reads via `readabilityHandler` rather than `FileHandle.bytes`:
            // the byte-by-byte `AsyncBytes` sequence terminates early / drops data on
            // large bursty CLI output, truncating a big tool-result line and losing
            // everything the agent streamed after it (the final answer + result event).
            let splitter = LineSplitter()
            handle.readabilityHandler = { fh in
                let chunk = fh.availableData
                if chunk.isEmpty {  // EOF
                    fh.readabilityHandler = nil
                    if let rest = splitter.drain() { continuation.yield(rest) }
                    continuation.finish()
                    return
                }
                for line in splitter.push(chunk) { continuation.yield(line) }
            }
            continuation.onTermination = { _ in handle.readabilityHandler = nil }
        }
    }

    func writeLine(_ data: Data) throws {
        stdinLock.lock()
        defer { stdinLock.unlock() }
        guard let stdinHandle, !stdinClosed else { throw BackendError.notRunning }
        var framed = data
        framed.append(0x0A)
        try stdinHandle.write(contentsOf: framed)
    }

    func closeStdin() {
        stdinLock.lock()
        defer { stdinLock.unlock() }
        guard !stdinClosed else { return }
        stdinClosed = true
        try? stdinHandle?.close()
    }

    func sendSignal(_ sig: Int32) {
        let pid = process.processIdentifier
        guard pid > 0 else { return }
        kill(pid, sig)
    }

    func wait() async -> Int32 {
        await exitTask.value
    }
}

/// Accumulates raw stdout/stderr bytes and splits them into newline-delimited lines,
/// holding any partial trailing line across chunk boundaries. `readabilityHandler`
/// fires serially, but it's locked anyway so a final drain can't race the last chunk.
///
/// Module-internal (not file-private) so the containerization adapter
/// (`ContainerizedProcessHandle`) can reuse the exact same NDJSON line framing the host
/// `ChildProcess` uses — guest stdio arrives as raw byte chunks over vsock and must be split
/// identically for the backend decode loop to behave the same in-container as on the host.
final class LineSplitter: @unchecked Sendable {
    private let lock = NSLock()
    private var buffer = Data()

    /// Append a chunk; return every complete line it produced (newline stripped, CR-trimmed).
    func push(_ chunk: Data) -> [Data] {
        lock.lock(); defer { lock.unlock() }
        buffer.append(chunk)
        var lines: [Data] = []
        var start = buffer.startIndex
        while let newline = buffer[start...].firstIndex(of: 0x0A) {
            var line = Data(buffer[start..<newline])
            if line.last == 0x0D { line.removeLast() }
            lines.append(line)
            start = buffer.index(after: newline)
        }
        // Rebase the unconsumed remainder so indices stay zero-based.
        buffer = start == buffer.startIndex ? buffer : Data(buffer[start...])
        return lines
    }

    /// Any trailing bytes with no terminating newline (flushed at EOF).
    func drain() -> Data? {
        lock.lock(); defer { lock.unlock() }
        guard !buffer.isEmpty else { return nil }
        defer { buffer.removeAll() }
        return buffer
    }
}
