import Darwin
import Foundation
import VMAgentCore

extension ConnectionHandler {
    /// Handle a streaming **`exec`** request (docs/MACOS_VM_NATIVE_AGENT.md §4.4): spawn the command and
    /// pump its stdio over this connection as base64 NDJSON frames until it exits, forwarding any
    /// host→guest `stdin`/`signal`/`eof` frames to the child. Owns ``fd`` for the whole run; the caller
    /// closes the connection afterwards (exec connections are single-use).
    ///
    /// The command is run under **`/bin/zsh -c`** — the guest account's login shell, invoked exactly as
    /// sshd used to invoke it, so `/etc/zshenv` (where provisioning appends the toolchain PATH) is
    /// sourced and `xcodebuild`/`brew`/`node` resolve. The host has already composed any working
    /// directory + env exports into the script.
    func runExec(requestLine: Data) {
        guard
            let parsed = try? JSONSerialization.jsonObject(with: requestLine),
            let request = parsed as? [String: Any],
            let command = request[AgentWire.Exec.commandKey] as? String
        else {
            sendExecError("exec requires a \"\(AgentWire.Exec.commandKey)\" string")
            return
        }

        let process = Process()
        process.executableURL = URL(fileURLWithPath: "/bin/zsh")
        process.arguments = ["-c", command]
        let outPipe = Pipe(), errPipe = Pipe(), inPipe = Pipe()
        process.standardOutput = outPipe
        process.standardError = errPipe
        process.standardInput = inPipe

        // Frame writes come from three places (both pipe-reader queues and this thread's exit frame),
        // so serialize them onto the one fd.
        let writeLock = NSLock()
        func send(_ frame: [String: Any]) {
            guard let data = try? JSONSerialization.data(withJSONObject: frame) else { return }
            var line = data
            line.append(0x0A)
            writeLock.lock()
            _ = writeAll(line)
            writeLock.unlock()
        }

        // Shared with the reader queues + terminationHandler; guarded because those run off-thread.
        let stateLock = NSLock()
        var stdoutOpen = true, stderrOpen = true, exited = false, exitCode: Int32 = 0

        outPipe.fileHandleForReading.readabilityHandler = { handle in
            let chunk = handle.availableData
            if chunk.isEmpty {
                stateLock.lock(); stdoutOpen = false; stateLock.unlock()
                handle.readabilityHandler = nil
            } else {
                send([AgentWire.Exec.tagKey: AgentWire.Exec.stdout,
                      AgentWire.Exec.dataKey: chunk.base64EncodedString()])
            }
        }
        errPipe.fileHandleForReading.readabilityHandler = { handle in
            let chunk = handle.availableData
            if chunk.isEmpty {
                stateLock.lock(); stderrOpen = false; stateLock.unlock()
                handle.readabilityHandler = nil
            } else {
                send([AgentWire.Exec.tagKey: AgentWire.Exec.stderr,
                      AgentWire.Exec.dataKey: chunk.base64EncodedString()])
            }
        }
        process.terminationHandler = { proc in
            stateLock.lock(); exited = true; exitCode = proc.terminationStatus; stateLock.unlock()
        }

        do {
            try process.run()
        } catch {
            outPipe.fileHandleForReading.readabilityHandler = nil
            errPipe.fileHandleForReading.readabilityHandler = nil
            sendExecError("could not launch the command: \(error)")
            return
        }

        var stdinClosed = false
        func closeStdin() {
            guard !stdinClosed else { return }
            try? inPipe.fileHandleForWriting.close()
            stdinClosed = true
        }
        func handleHostFrame(_ line: Data) {
            guard
                let parsed = try? JSONSerialization.jsonObject(with: line),
                let frame = parsed as? [String: Any],
                let tag = frame[AgentWire.Exec.tagKey] as? String
            else { return }
            switch tag {
            case AgentWire.Exec.stdin:
                if let b64 = frame[AgentWire.Exec.dataKey] as? String,
                    let bytes = Data(base64Encoded: b64), !stdinClosed
                {
                    try? inPipe.fileHandleForWriting.write(contentsOf: bytes)
                }
            case AgentWire.Exec.stdinEOF:
                closeStdin()
            case AgentWire.Exec.signal:
                if let n = frame[AgentWire.Exec.signalKey] as? Int {
                    kill(process.processIdentifier, Int32(n))
                }
            default:
                break
            }
        }

        // Any host→guest frames pipelined right after the request line are already buffered.
        while let line = lines.nextLine() {
            if !line.isEmpty { handleHostFrame(line) }
        }

        // Pump the connection until the process has exited AND both output pipes have drained. A short
        // poll tick lets us notice the exit even when the host sends no more input frames.
        var scratch = [UInt8](repeating: 0, count: 64 * 1024)
        while true {
            stateLock.lock()
            let finished = exited && !stdoutOpen && !stderrOpen
            stateLock.unlock()
            if finished { break }

            var pfd = pollfd(fd: fd, events: Int16(POLLIN), revents: 0)
            let ready = poll(&pfd, 1, 200)
            if ready > 0, (pfd.revents & Int16(POLLIN)) != 0 {
                let n = read(fd, &scratch, scratch.count)
                if n > 0 {
                    lines.append(Data(bytes: scratch, count: n))
                    while let line = lines.nextLine() {
                        if !line.isEmpty { handleHostFrame(line) }
                    }
                } else if n == 0 {
                    // Host hung up — kill the command and stop streaming.
                    kill(process.processIdentifier, SIGKILL)
                    break
                } else if errno != EINTR && errno != EAGAIN {
                    kill(process.processIdentifier, SIGKILL)
                    break
                }
            }
        }

        closeStdin()
        outPipe.fileHandleForReading.readabilityHandler = nil
        errPipe.fileHandleForReading.readabilityHandler = nil
        stateLock.lock(); let code = exitCode; stateLock.unlock()
        send([AgentWire.Exec.tagKey: AgentWire.Exec.exit, AgentWire.Exec.codeKey: Int(code)])
    }

    /// Emit a single terminal `spawnError` frame (the exec never started).
    private func sendExecError(_ message: String) {
        guard
            let data = try? JSONSerialization.data(withJSONObject: [
                AgentWire.Exec.tagKey: AgentWire.Exec.spawnError,
                AgentWire.Exec.messageKey: message,
            ])
        else { return }
        var line = data
        line.append(0x0A)
        _ = writeAll(line)
    }
}
