import CVsock
import Darwin
import Foundation
import VMAgentCore

enum AgentError: Error, CustomStringConvertible {
    case posix(String, Int32)
    var description: String {
        switch self {
        case let .posix(what, errno): return "\(what) failed: \(String(cString: strerror(errno)))"
        }
    }
}

/// Timestamped line to stderr — the LaunchAgent routes it to the agent's log file.
func agentLog(_ message: String) {
    FileHandle.standardError.write(Data("[NucleicVMAgent] \(message)\n".utf8))
}

/// Owns the AF_VSOCK listening socket and the blocking accept loop. One `ConnectionHandler` thread
/// per accepted connection; in practice there is exactly one client (the host's `MacVMAgentClient`),
/// but a reconnect after a host restart must not require an agent restart.
final class AgentListener {
    private let fd: Int32

    init(port: UInt32) throws {
        let result = nucleic_vsock_listen(port, 4)
        guard result >= 0 else { throw AgentError.posix("vsock listen(:\(port))", -result) }
        fd = result
    }

    func run() {
        while true {
            let conn = accept(fd, nil, nil)
            if conn < 0 {
                if errno == EINTR { continue }
                agentLog("accept failed (errno \(errno)); retrying")
                Thread.sleep(forTimeInterval: 1)
                continue
            }
            agentLog("host connected (fd \(conn))")
            let handler = ConnectionHandler(fd: conn)
            Thread.detachNewThread { handler.run() }
        }
    }
}

/// One host connection: a blocking read loop that frames NDJSON requests, dispatches each op, and
/// writes the one-line JSON reply. Everything for this connection (including the AX element
/// registry) is confined to this thread — requests on a connection are strictly serial.
final class ConnectionHandler {
    /// The accepted vsock connection fd. Internal (not private) so the ``runExec`` streamer — which
    /// takes the connection over for the life of one command — can poll and write it directly.
    let fd: Int32
    /// NDJSON line framer for this connection. Internal so ``runExec`` can keep draining host→guest
    /// frames (stdin / signal) from the same buffer after the request line.
    var lines = LineSplitBuffer()
    /// AX element handles (`ref` → live `AXUIElement`) issued by the last dump on this connection.
    let axRegistry = AXRegistry()

    init(fd: Int32) {
        self.fd = fd
    }

    func run() {
        defer {
            close(fd)
            agentLog("host disconnected (fd \(fd))")
        }
        var scratch = [UInt8](repeating: 0, count: 64 * 1024)
        while true {
            let n = read(fd, &scratch, scratch.count)
            if n == 0 { return }  // EOF — host closed
            if n < 0 {
                if errno == EINTR { continue }
                return
            }
            lines.append(Data(bytes: scratch, count: n))
            while let line = lines.nextLine() {
                guard !line.isEmpty else { continue }
                // `exec` is the one streaming op: it takes over this connection for the life of the
                // command (bidirectional NDJSON frames, §6), then the connection is done — exec
                // connections are single-use, so return and let `defer` close the fd.
                if Self.isExecRequest(line) {
                    runExec(requestLine: line)
                    return
                }
                var reply = handle(requestLine: line)
                reply.append(0x0A)
                guard writeAll(reply) else { return }
            }
            // Runaway-line guard: a client that never sends a newline can't balloon memory.
            if lines.pendingBytes > 4 * 1024 * 1024 { return }
        }
    }

    /// Peek a framed request line for `"op":"exec"` without disturbing dispatch of every other op.
    private static func isExecRequest(_ line: Data) -> Bool {
        guard
            let parsed = try? JSONSerialization.jsonObject(with: line),
            let request = parsed as? [String: Any]
        else { return false }
        return (request["op"] as? String) == AgentWire.Exec.op
    }

    func writeAll(_ data: Data) -> Bool {
        var remaining = data
        while !remaining.isEmpty {
            let written = remaining.withUnsafeBytes { raw in
                write(fd, raw.baseAddress, raw.count)
            }
            if written < 0 {
                if errno == EINTR { continue }
                return false
            }
            remaining = remaining.dropFirst(written)
        }
        return true
    }

    // MARK: - Dispatch

    private func handle(requestLine: Data) -> Data {
        let reply = dispatch(requestLine: requestLine)
        if let data = try? JSONSerialization.data(withJSONObject: reply) { return data }
        return Data(#"{"ok":false,"error":"reply serialization failed"}"#.utf8)
    }

    private func dispatch(requestLine: Data) -> [String: Any] {
        guard
            let parsed = try? JSONSerialization.jsonObject(with: requestLine),
            let request = parsed as? [String: Any],
            let op = request["op"] as? String
        else {
            return AgentReply.err("malformed request (want one JSON object per line with an \"op\")")
        }
        // Each op is individually guarded so a framework-level surprise degrades to an error reply,
        // never a dead agent.
        switch op {
        case "ping": return opPing()
        case "screenshot": return opScreenshot(request)
        case "ax_dump": return opAXDump(request)
        case "ax_element_at": return opAXElementAt(request)
        case "ax_action": return opAXAction(request)
        case "ax_set_value": return opAXSetValue(request)
        case "ax_focus": return opAXFocus(request)
        case "click": return opClick(request)
        case "move": return opMove(request)
        case "drag": return opDrag(request)
        case "type": return opType(request)
        case "key": return opKey(request)
        case "scroll": return opScroll(request)
        case "cursor_position": return opCursorPosition()
        case "launch_app": return opLaunchApp(request)
        default:
            return AgentReply.err("unknown op '\(op)' (agent protocol v\(AgentWire.version))")
        }
    }
}

/// Reply-shape helpers (§4): `{ok:true, …}` or `{ok:false, error}`.
enum AgentReply {
    static func ok(_ fields: [String: Any] = [:]) -> [String: Any] {
        var reply = fields
        reply["ok"] = true
        return reply
    }

    static func err(_ message: String) -> [String: Any] {
        ["ok": false, "error": message]
    }
}
