import Foundation
import NucleicProtocol
// Unconditional: `VZVirtioSocketConnection` (unlike the arm64-only `VZMac*` guest APIs) exists on
// both architectures, and the client actor must be a real type wherever the engine compiles.
import Virtualization

/// Wire constants + framing for the host ⇄ in-guest-agent protocol
/// (docs/MACOS_VM_NATIVE_AGENT.md §3–§4): NDJSON over vsock — one JSON request object per line
/// (carrying an `op`), one JSON reply per line (carrying `ok` + op fields, or `ok:false, error`).
///
/// The guest half lives in `guest/NucleicVMAgent` (its own package — the agent runs on macOS 14+
/// guests while this package's floor is 26), which mirrors these constants in `VMAgentCore.AgentWire`;
/// keep the two in lockstep.
enum MacVMAgentWire {
    /// The Nucleic-reserved vsock port the guest agent listens on.
    static let port: UInt32 = 2035
    /// The Linux **a11y (AT-SPI) semantic agent** vsock port (docs/LINUX_VM_SEMANTIC_AGENT.md §9 Q2):
    /// a second port so the static-C `ping`/`exec` agent on ``port`` stays untouched. Same NDJSON wire.
    static let a11yPort: UInt32 = 2036
    /// Protocol version this client speaks; the agent reports its own in the `ping` reply.
    /// v2 adds the streaming ``Exec`` op — the vsock replacement for the retired SSH exec path.
    /// Mirror of `VMAgentCore.AgentWire.version`; keep the two in lockstep.
    static let version = 2

    /// The streaming **`exec`** sub-protocol (docs/MACOS_VM_NATIVE_AGENT.md §4.4), mirror of
    /// `VMAgentCore.AgentWire.Exec`. After the single `{"op":"exec","command":…}` request line, the
    /// (dedicated) connection carries one-JSON-object-per-line *frames* tagged by ``tagKey``. Payload
    /// bytes are base64 in ``dataKey``. See ``MacVMExecChannel``.
    enum Exec {
        static let op = "exec"
        static let commandKey = "command"
        static let tagKey = "t"
        // guest → host
        static let stdout = "o"
        static let stderr = "e"
        static let exit = "x"
        static let spawnError = "err"
        // host → guest
        static let stdin = "i"
        static let stdinEOF = "eof"
        static let signal = "sig"
        // frame field keys
        static let dataKey = "d"
        static let codeKey = "code"
        static let messageKey = "m"
        static let signalKey = "n"
    }
    /// Default per-request deadline. Generous enough for a full-display SCK screenshot or a deep
    /// AX walk; a hung agent surfaces as `.timeout` instead of wedging the tool call.
    static let requestTimeout: TimeInterval = 20

    /// Encode one request line (`{"op": …, …}\n`). Pure, unit-tested.
    static func requestLine(op: String, fields: [String: JSONValue] = [:]) -> Data {
        var object = fields
        object["op"] = .string(op)
        return requestLineRaw(object)
    }

    /// Encode one NDJSON line from a raw object (no forced `op`) — used for the streaming exec frames
    /// (`{"t":…}`), which carry a tag rather than an op. Pure.
    static func requestLineRaw(_ fields: [String: JSONValue]) -> Data {
        var line = (try? JSONValue.object(fields).encodedData()) ?? Data("{}".utf8)
        line.append(0x0A)
        return line
    }

    /// Decode one reply line into its JSON object. Pure, unit-tested.
    static func parseReply(_ line: Data) throws -> JSONValue {
        guard let reply = try? JSONValue(parsing: line), reply.objectValue != nil else {
            throw MacVMAgentError.protocolError(
                "malformed agent reply: \(String(decoding: line.prefix(200), as: UTF8.self))")
        }
        return reply
    }
}

/// Errors from the native-agent channel. Deliberately separate from ``MacVMError``: these are
/// routing signals — the computer-use path catches them and falls back to SSH — not user-facing
/// VM-lifecycle failures.
enum MacVMAgentError: Error, CustomStringConvertible {
    /// The VM was built without a vsock device (pre-agent configuration still booted).
    case noSocketDevice
    /// The vsock connection died or the fd errored; the client is unusable.
    case ioFailed(String)
    /// No reply within the deadline (agent hung, or nothing listening behind the connect).
    case timeout
    /// The reply wasn't the protocol's shape.
    case protocolError(String)
    /// The agent executed the op and reported failure (`ok:false`) — transport is healthy.
    case agentError(String)

    var description: String {
        switch self {
        case .noSocketDevice: return "the VM has no vsock device"
        case .ioFailed(let why): return "agent connection failed: \(why)"
        case .timeout: return "the in-guest agent did not reply in time"
        case .protocolError(let why): return why
        case .agentError(let message): return message
        }
    }

    /// Whether the underlying connection should be discarded (vs. an op-level `agentError`, after
    /// which the channel remains perfectly usable).
    var isTransportFailure: Bool {
        switch self {
        case .agentError: return false
        default: return true
        }
    }
}

/// One live NDJSON request/reply channel to the in-guest `NucleicVMAgent`, over the vsock
/// connection ``MacVMInstance/connectAgent(port:)`` opened. The actor serializes requests (the
/// protocol is strictly one-reply-per-request), owns the read buffer, and keeps the
/// `VZVirtioSocketConnection` alive while I/O is in flight (its fd closes with it).
///
/// I/O is nonblocking `read`/`write` on the connection's fd with `Task.sleep` backoff — no thread
/// is ever parked on a guest that stops answering; the deadline turns it into `.timeout`.
actor MacVMAgentClient {
    private let connection: UncheckedSendableBox<VZVirtioSocketConnection>
    private let fd: Int32
    private var buffer = Data()
    private var closed = false

    init(connection: UncheckedSendableBox<VZVirtioSocketConnection>) {
        self.connection = connection
        self.fd = connection.value.fileDescriptor
        // Nonblocking so the actor yields (rather than blocks a cooperative thread) while waiting.
        let flags = fcntl(fd, F_GETFL, 0)
        _ = fcntl(fd, F_SETFL, flags | O_NONBLOCK)
    }

    /// Send one op and await its reply object. Throws ``MacVMAgentError/agentError(_:)`` when the
    /// agent answers `ok:false`, a transport error when the channel itself broke.
    func request(
        op: String, fields: [String: JSONValue] = [:],
        timeout: TimeInterval = MacVMAgentWire.requestTimeout
    ) async throws -> JSONValue {
        guard !closed, fd >= 0 else { throw MacVMAgentError.ioFailed("connection closed") }
        let deadline = Date().addingTimeInterval(timeout)
        try await writeAll(MacVMAgentWire.requestLine(op: op, fields: fields), deadline: deadline)
        let line = try await readLine(deadline: deadline)
        let reply = try MacVMAgentWire.parseReply(line)
        guard reply["ok"]?.boolValue == true else {
            throw MacVMAgentError.agentError(
                reply["error"]?.stringValue ?? "the agent reported an unspecified failure")
        }
        return reply
    }

    /// Close the vsock connection. Idempotent; subsequent requests throw.
    func close() {
        guard !closed else { return }
        closed = true
        connection.value.close()
    }

    // MARK: - Nonblocking fd I/O

    private func writeAll(_ data: Data, deadline: Date) async throws {
        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 {
                guard Date() < deadline else { throw MacVMAgentError.timeout }
                try? await Task.sleep(nanoseconds: 20_000_000)
            } else {
                closed = true
                throw MacVMAgentError.ioFailed("write errno \(errno)")
            }
        }
    }

    private func readLine(deadline: Date) async throws -> Data {
        var scratch = [UInt8](repeating: 0, count: 256 * 1024)
        while true {
            if let newline = buffer.firstIndex(of: 0x0A) {
                let line = buffer.subdata(in: buffer.startIndex..<newline)
                buffer.removeSubrange(buffer.startIndex...newline)
                return line
            }
            // A screenshot reply is a few MB of base64 — cap far above that, but cap.
            if buffer.count > 32 * 1024 * 1024 {
                closed = true
                throw MacVMAgentError.protocolError("agent reply exceeded 32 MB without a newline")
            }
            let count = read(fd, &scratch, scratch.count)
            if count > 0 {
                buffer.append(contentsOf: scratch[0..<count])
            } else if count == 0 {
                closed = true
                throw MacVMAgentError.ioFailed("the agent closed the connection")
            } else if errno == EAGAIN || errno == EINTR {
                guard Date() < deadline else { throw MacVMAgentError.timeout }
                try? await Task.sleep(nanoseconds: 25_000_000)
            } else {
                closed = true
                throw MacVMAgentError.ioFailed("read errno \(errno)")
            }
        }
    }
}

#if arch(arm64)
extension MacVMInstance {
    /// Open a vsock connection to `port` in this VM's guest. Runs the connect on the VM's serial
    /// queue (all VZ calls must); resolves once the guest's listener accepts. If nothing in the
    /// guest listens on the port (no agent baked in / not yet launched), the framework fails the
    /// connect — the engine's probe turns that into "agent unavailable", never an error the user sees.
    func connectAgent(port: UInt32) async throws -> UncheckedSendableBox<VZVirtioSocketConnection> {
        try await withCheckedThrowingContinuation { cont in
            queue.async { [self] in
                guard let device = vm.socketDevices.first as? VZVirtioSocketDevice else {
                    cont.resume(throwing: MacVMAgentError.noSocketDevice)
                    return
                }
                device.connect(toPort: port) { result in
                    switch result {
                    case .success(let connection):
                        cont.resume(returning: UncheckedSendableBox(value: connection))
                    case .failure(let error):
                        cont.resume(throwing: error)
                    }
                }
            }
        }
    }
}
#endif

extension MacVMEngine {
    /// The native in-guest agent channel for a live VM, or `nil` when the VM has none (older base
    /// image, agent not running, or vsock connect unsupported) — the caller then uses the SSH path.
    ///
    /// First call per VM probes: vsock connect + `ping`, both under a short deadline. The result is
    /// cached on the ``LiveVM`` (see ``MacVMAgentState``); a concurrent duplicate probe is harmless
    /// (one extra short-lived connection — the agent accepts several).
    func agentClient(name: String) async -> MacVMAgentClient? {
        #if arch(arm64)
        guard let entry = live[name] else { return nil }
        switch entry.agent {
        case .available(let client): return client
        case .unavailable: return nil
        case .unprobed: break
        }
        do {
            let instance = entry.instance
            let box = try await withAgentTimeout(seconds: 5) {
                try await instance.connectAgent(port: MacVMAgentWire.port)
            }
            let client = MacVMAgentClient(connection: box)
            let pong = try await client.request(op: "ping", timeout: 5)
            // Readiness is informational (each op re-checks its own TCC grant and reports
            // precisely); version is the only hard gate.
            guard let version = pong["version"]?.intValue, version == MacVMAgentWire.version else {
                await client.close()
                throw MacVMAgentError.protocolError(
                    "agent protocol v\(pong["version"]?.intValue ?? 0) ≠ host v\(MacVMAgentWire.version)")
            }
            live[name]?.agent = .available(client)
            return client
        } catch {
            live[name]?.agent = .unavailable
            return nil
        }
        #else
        return nil
        #endif
    }

    /// The Linux **a11y (AT-SPI) semantic agent** channel for a live Linux VM, or `nil` when the VM
    /// has none (base built without the agent, or it isn't up in the session yet) — the caller then
    /// falls back to screenshot + pixel actions. Mirrors ``agentClient(name:)`` but on the a11y vsock
    /// port and cached separately on ``LiveVM/a11yAgent`` (the C `ping`/`exec` agent on the main port
    /// is a different connection).
    func linuxA11yClient(name: String) async -> MacVMAgentClient? {
        #if arch(arm64)
        guard let entry = live[name] else { return nil }
        switch entry.a11yAgent {
        case .available(let client): return client
        case .unavailable: return nil
        case .unprobed: break
        }
        do {
            let instance = entry.instance
            let box = try await withAgentTimeout(seconds: 5) {
                try await instance.connectAgent(port: MacVMAgentWire.a11yPort)
            }
            let client = MacVMAgentClient(connection: box)
            let pong = try await client.request(op: "ping", timeout: 5)
            guard let version = pong["version"]?.intValue, version == MacVMAgentWire.version else {
                await client.close()
                throw MacVMAgentError.protocolError("a11y agent protocol mismatch")
            }
            live[name]?.a11yAgent = .available(client)
            return client
        } catch {
            live[name]?.a11yAgent = .unavailable
            return nil
        }
        #else
        return nil
        #endif
    }

    /// Discard a VM's agent channel after a transport failure so the next action re-probes once
    /// (the guest LaunchAgent's `KeepAlive` may have already restarted the agent).
    func discardAgentClient(name: String, _ client: MacVMAgentClient) async {
        #if arch(arm64)
        await client.close()
        if live[name] != nil {
            if case .available(let a) = live[name]?.agent, a === client { live[name]?.agent = .unprobed }
            if case .available(let a) = live[name]?.a11yAgent, a === client { live[name]?.a11yAgent = .unprobed }
        }
        #endif
    }
}

/// Race `body` against a wall-clock deadline. The vsock `connect` has no timeout of its own and
/// Apple documents the no-listener case only as "does nothing" — this guarantees the agent probe
/// resolves either way.
func withAgentTimeout<T: Sendable>(
    seconds: TimeInterval, _ body: @escaping @Sendable () async throws -> T
) async throws -> T {
    try await withThrowingTaskGroup(of: T.self) { group in
        group.addTask { try await body() }
        group.addTask {
            try await Task.sleep(nanoseconds: UInt64(seconds * 1_000_000_000))
            throw MacVMAgentError.timeout
        }
        // One of the two always completes; cancel the loser.
        let winner = try await group.next()!
        group.cancelAll()
        return winner
    }
}
