import Foundation
import NucleicProtocol
import Virtualization

#if arch(arm64)

/// A ``ProcessHandle`` whose remote process runs **inside a macOS guest, driven entirely over vsock**
/// (docs/MACOS_VM_NATIVE_AGENT.md §4.4) — the direct replacement for the retired `ssh` exec handle.
///
/// It owns one **dedicated** `VZVirtioSocketConnection` to the in-guest agent (port 2035). On init it
/// sends a single `{"op":"exec","command":…}` line; from then on the connection carries the streaming
/// exec frames: base64 `stdout`/`stderr` chunks and a final `exit`. The class re-splits those chunks
/// into newline-delimited `Data` lines so the backend decode loop can't tell a guest exec from a host
/// or containerized one — the same trick §2.1 of docs/MACOS_VM.md called out for the SSH handle.
///
/// `stdin` (``writeLine``), `closeStdin`, and `sendSignal` travel the other way as host→guest frames.
/// `@unchecked Sendable`: all mutable state is guarded by `lock`, and the fd is full-duplex (the
/// reader task reads while callers write control frames).
final class MacVMExecChannel: ProcessHandle, @unchecked Sendable {
    let stdoutLines: AsyncThrowingStream<Data, Error>
    let stderrLines: AsyncThrowingStream<Data, Error>
    private let stdoutCont: AsyncThrowingStream<Data, Error>.Continuation
    private let stderrCont: AsyncThrowingStream<Data, Error>.Continuation

    private let connection: UncheckedSendableBox<VZVirtioSocketConnection>
    private let fd: Int32
    private let lock = NSLock()

    // Reader-side line accumulators (touched only on the reader task).
    private var stdoutAcc = Data()
    private var stderrAcc = Data()

    /// Most recent guest stdout/stderr chunk — the sign of life the stall monitor reads. The remote
    /// pid lives in the VM, so ``supportsHostCPUSampling`` stays `false` (output-idle only).
    private let activity = ActivityClock()
    var lastActivityNanos: UInt64 { activity.lastNanos }

    // Guarded by `lock`.
    private var finished = false
    private var exitCode: Int32?
    private var exitWaiters: [CheckedContinuation<Int32, Never>] = []

    /// Open an exec channel over `connection`, immediately dispatching `command` (a shell script the
    /// engine has already composed with env exports + working directory). The request carries no
    /// shell choice: the guest runs **nash, and only nash** (docs/NASH.md §7.5 — locked; a nash-less
    /// guest refuses with a `spawnError` naming the fix). The reader task starts pumping frames
    /// right away.
    init(connection: UncheckedSendableBox<VZVirtioSocketConnection>, command: String) {
        self.connection = connection
        self.fd = connection.value.fileDescriptor
        var out: AsyncThrowingStream<Data, Error>.Continuation!
        self.stdoutLines = AsyncThrowingStream { out = $0 }
        self.stdoutCont = out
        var err: AsyncThrowingStream<Data, Error>.Continuation!
        self.stderrLines = AsyncThrowingStream { err = $0 }
        self.stderrCont = err

        // Nonblocking so the reader task yields (never parks a cooperative thread) between chunks.
        let flags = fcntl(fd, F_GETFL, 0)
        _ = fcntl(fd, F_SETFL, flags | O_NONBLOCK)

        let request = MacVMAgentWire.requestLine(
            op: MacVMAgentWire.Exec.op,
            fields: [MacVMAgentWire.Exec.commandKey: .string(command)])
        writeFrame(request)

        Task.detached { [weak self] in await self?.readLoop() }
    }

    // MARK: - ProcessHandle

    /// Informational only — real signalling rides ``sendSignal`` frames, not a host pid.
    var processID: Int32 { fd }

    /// Send one line to the remote process's stdin as an `stdin` frame (newline appended, matching the
    /// host pipe's line-oriented `writeLine`).
    func writeLine(_ data: Data) throws {
        var payload = data
        payload.append(0x0A)
        sendControl(tag: MacVMAgentWire.Exec.stdin, dataKey: payload.base64EncodedString())
    }

    /// Stream one **raw binary** stdin payload as a single `stdin` frame — for bulk transfers (an app
    /// tarball) where ``writeLine``'s appended `0x0A` would corrupt the byte stream. Unlike the
    /// control-frame writers, it does NOT silently give up after a fixed spin budget: it blocks against
    /// `deadline`, so sustained guest backpressure (a slow `tar` draining stdin) throttles the sender
    /// instead of dropping a chunk. Returns `false` if the channel is torn down or the deadline passes.
    func writeStdinRaw(_ data: Data, deadline: Date) -> Bool {
        let frame = MacVMAgentWire.requestLineRaw([
            MacVMAgentWire.Exec.tagKey: .string(MacVMAgentWire.Exec.stdin),
            MacVMAgentWire.Exec.dataKey: .string(data.base64EncodedString()),
        ])
        return writeFrameBlocking(frame, deadline: deadline)
    }

    func closeStdin() {
        sendControl(tag: MacVMAgentWire.Exec.stdinEOF, dataKey: nil)
    }

    func sendSignal(_ sig: Int32) {
        let frame = MacVMAgentWire.requestLineRaw([
            MacVMAgentWire.Exec.tagKey: .string(MacVMAgentWire.Exec.signal),
            MacVMAgentWire.Exec.signalKey: .number(Double(sig)),
        ])
        writeFrame(frame)
    }

    func wait() async -> Int32 {
        await withCheckedContinuation { (cont: CheckedContinuation<Int32, Never>) in
            let alreadyExited: Int32? = lock.withLock {
                if let code = exitCode { return code }
                exitWaiters.append(cont)
                return nil
            }
            if let alreadyExited { cont.resume(returning: alreadyExited) }
        }
    }

    /// Finish the line streams immediately (the caller is done reading), independent of the guest.
    /// Mirrors ``ContainerizedProcessHandle`` — a killed exec's streams shouldn't outlive interest.
    func forceCloseStreams() {
        finish(code: nil)
    }

    // MARK: - Reader

    private func readLoop() async {
        var frameBuf = Data()
        var scratch = [UInt8](repeating: 0, count: 256 * 1024)
        while true {
            let count = read(fd, &scratch, scratch.count)
            if count > 0 {
                frameBuf.append(contentsOf: scratch[0..<count])
                // A pathological producer without newlines can't balloon host memory.
                if frameBuf.count > 64 * 1024 * 1024 { break }
                while let nl = frameBuf.firstIndex(of: 0x0A) {
                    let frame = frameBuf.subdata(in: frameBuf.startIndex..<nl)
                    frameBuf.removeSubrange(frameBuf.startIndex...nl)
                    if handleFrame(frame) { return }  // exit frame → streams already finished
                }
            } else if count == 0 {
                break  // guest closed the connection without an exit frame
            } else if errno == EAGAIN || errno == EINTR {
                if lock.withLock({ finished }) { return }
                try? await Task.sleep(nanoseconds: 15_000_000)
            } else {
                break  // fd error
            }
        }
        // Fell out without an `exit` frame: transport died. Surface a nonzero code.
        finish(code: exitCode ?? -1)
    }

    /// Decode one guest→host frame. Returns `true` once the terminal (`exit`/`spawnError`) frame has
    /// been handled and the streams finished.
    private func handleFrame(_ data: Data) -> Bool {
        guard
            let obj = (try? JSONSerialization.jsonObject(with: data)) as? [String: Any],
            let tag = obj[MacVMAgentWire.Exec.tagKey] as? String
        else { return false }
        switch tag {
        case MacVMAgentWire.Exec.stdout:
            if let b64 = obj[MacVMAgentWire.Exec.dataKey] as? String,
                let bytes = Data(base64Encoded: b64)
            {
                emitLines(bytes, into: stdoutCont, accumulator: &stdoutAcc)
            }
            return false
        case MacVMAgentWire.Exec.stderr:
            if let b64 = obj[MacVMAgentWire.Exec.dataKey] as? String,
                let bytes = Data(base64Encoded: b64)
            {
                emitLines(bytes, into: stderrCont, accumulator: &stderrAcc)
            }
            return false
        case MacVMAgentWire.Exec.spawnError:
            let message = obj[MacVMAgentWire.Exec.messageKey] as? String ?? "the guest could not run the command"
            emitLines(Data(message.utf8), into: stderrCont, accumulator: &stderrAcc)
            finish(code: 127)
            return true
        case MacVMAgentWire.Exec.exit:
            let code = obj[MacVMAgentWire.Exec.codeKey] as? Int ?? 0
            finish(code: Int32(code))
            return true
        default:
            return false
        }
    }

    /// Append `bytes` to `accumulator` and yield every complete newline-terminated line to `cont`
    /// (newline stripped), holding any partial trailing line for the next chunk.
    private func emitLines(
        _ bytes: Data, into cont: AsyncThrowingStream<Data, Error>.Continuation,
        accumulator: inout Data
    ) {
        activity.bump()
        accumulator.append(bytes)
        while let nl = accumulator.firstIndex(of: 0x0A) {
            let line = accumulator.subdata(in: accumulator.startIndex..<nl)
            accumulator.removeSubrange(accumulator.startIndex...nl)
            cont.yield(line)
        }
    }

    // MARK: - Teardown

    /// Finish both streams (flushing any trailing partial line), record the exit code, wake every
    /// `wait()`, and close the connection. Idempotent.
    private func finish(code: Int32?) {
        lock.lock()
        if finished {
            lock.unlock()
            return
        }
        finished = true
        let resolved = code ?? exitCode ?? -1
        exitCode = resolved
        let waiters = exitWaiters
        exitWaiters.removeAll()
        lock.unlock()

        if !stdoutAcc.isEmpty { stdoutCont.yield(stdoutAcc); stdoutAcc.removeAll() }
        if !stderrAcc.isEmpty { stderrCont.yield(stderrAcc); stderrAcc.removeAll() }
        stdoutCont.finish()
        stderrCont.finish()
        for waiter in waiters { waiter.resume(returning: resolved) }
        connection.value.close()
    }

    // MARK: - Writes (host → guest)

    private func sendControl(tag: String, dataKey: String?) {
        var fields: [String: JSONValue] = [MacVMAgentWire.Exec.tagKey: .string(tag)]
        if let dataKey { fields[MacVMAgentWire.Exec.dataKey] = .string(dataKey) }
        writeFrame(MacVMAgentWire.requestLineRaw(fields))
    }

    /// Fully write one frame, blocking on backpressure until done or `deadline` passes — the reliable
    /// counterpart to ``writeFrame`` for bulk stdin (``writeStdinRaw``), which must never drop a chunk.
    /// Returns whether the whole frame was written.
    private func writeFrameBlocking(_ data: Data, deadline: Date) -> Bool {
        lock.lock()
        let done = finished
        lock.unlock()
        if done { return false }
        var remaining = data
        while !remaining.isEmpty {
            let written = remaining.withUnsafeBytes { raw in write(fd, raw.baseAddress, raw.count) }
            if written > 0 {
                remaining = remaining.dropFirst(written)
            } else if written < 0, errno == EAGAIN || errno == EINTR {
                if Date() >= deadline { return false }
                lock.lock()
                let torn = finished
                lock.unlock()
                if torn { return false }
                usleep(2000)
            } else {
                return false  // fd error — the read loop surfaces the teardown
            }
        }
        return true
    }

    /// Blocking-ish write of a whole control frame (tiny), serialized against other writers. Spins on
    /// EAGAIN since the fd is nonblocking; drops the frame if the connection has been torn down.
    private func writeFrame(_ data: Data) {
        lock.lock()
        let done = finished
        lock.unlock()
        if done { return }
        var remaining = data
        var spins = 0
        while !remaining.isEmpty {
            let written = remaining.withUnsafeBytes { raw in write(fd, raw.baseAddress, raw.count) }
            if written > 0 {
                remaining = remaining.dropFirst(written)
            } else if written < 0, errno == EAGAIN || errno == EINTR {
                spins += 1
                if spins > 1000 { return }  // ~1s of a wedged fd — give up on this frame
                usleep(1000)
            } else {
                return  // fd error — the read loop will surface the teardown
            }
        }
    }
}

#endif
