import Darwin
import Foundation

/// The streaming **`exec`** engine (docs/MACOS_VM_NATIVE_AGENT.md §4.4, docs/NASH.md §7.5): spawns the
/// requested command under the guest's shell and pumps its stdio over one connection fd as base64
/// NDJSON frames until it exits, forwarding host→guest `stdin`/`signal`/`eof` frames to the child.
///
/// Lives in `VMAgentCore` (not the executable target) so `VMAgentCoreTests` can drive a real session
/// over a `socketpair` — the transport is just an fd, so the tests exercise the exact code the vsock
/// connection runs. Single-threaded by construction: one `poll(2)` loop over the connection and both
/// output pipes, so there are no handler races and strict concurrency has nothing to object to.
///
/// Two guarantees this type exists to enforce — both regressions the Foundation-`Process` + pipe-EOF
/// implementation it replaced could not make (the "mac_vm_exec never returns" wedge, NASH.md §7.5):
///
/// 1. **The exit frame always arrives once the shell exits.** The old loop waited for the process AND
///    EOF on both output pipes. `zsh -c` mostly hid the difference by `exec`ing a single command in
///    place; nash (brush) always forks its children — so a backgrounded child, or any descendant
///    surviving a kill of the shell's pid, inherited the pipe write-ends and held EOF off forever, and
///    the host never got its exit frame. Now: once the child is reaped, remaining output gets a short
///    ``pipeGraceMillis`` to drain, then the exit frame is sent regardless. A survivor (e.g. a
///    deliberately `nohup`ed server) keeps running — it just can't wedge the tool call.
/// 2. **Signals reach the whole command tree.** The shell is spawned as its own process-group leader
///    (`POSIX_SPAWN_SETPGROUP`), and `signal` frames / host-hangup kills target the group
///    (`kill(-pid, …)`, falling back to the pid). Killing only the shell's pid was how orphans got
///    made in the first place.
public final class ExecSession {
    /// Where provisioning installs the universal nash binary (docs/NASH.md §7.5, §8).
    public static let defaultNashPath = "/usr/local/bin/nash"
    /// How long, after the shell itself has exited, to keep draining output pipes still held open by
    /// surviving descendants before sending the exit frame anyway. Long enough for a normal child's
    /// final buffered output to arrive; short enough that a survivor can't wedge the tool call.
    public static let defaultPipeGraceMillis: Int32 = 2000

    private let fd: Int32
    private var lines: LineSplitBuffer
    private let nashPath: String
    private let graceMillis: Int32

    public init(
        fd: Int32, buffered: LineSplitBuffer = LineSplitBuffer(),
        nashPath: String = ExecSession.defaultNashPath,
        pipeGraceMillis: Int32 = ExecSession.defaultPipeGraceMillis
    ) {
        self.fd = fd
        self.lines = buffered
        self.nashPath = nashPath
        self.graceMillis = pipeGraceMillis
        // A frame write racing the peer's hangup must fail with EPIPE, not raise SIGPIPE — the agent
        // ignores SIGPIPE process-wide, but this type must not depend on its host process doing so
        // (the tests are their own process). Darwin's per-fd F_SETNOSIGPIPE covers pipes and sockets.
        _ = fcntl(fd, F_SETNOSIGPIPE, 1)
    }

    // MARK: - Shell selection (docs/NASH.md §7.5)

    /// Pick the shell for one exec: **nash, and only nash** (locked decision — nash is the sole
    /// agent-accessible shell on every virtualized surface, so its observation has no gaps). There
    /// is no zsh fallback and no per-request override: a guest without nash installed refuses the
    /// exec (`nil` — the caller replies with an actionable `spawnError`) rather than silently
    /// degrading the observability boundary to zsh; provisioning hard-requires nash for exactly
    /// that reason. The one exception is the §4.1 operator break-glass — `disable=1` in the trusted
    /// policy file (``NashPolicy/path``), which an operator writes with `sudo` inside the guest.
    /// It used to be `NUCLEIC_NASH_DISABLE=1` in the agent daemon's environment; that channel is
    /// gone on purpose (docs/NASH.md §4.3), because nash itself no longer honors the variable and
    /// leaving it here would mean the top-level shell degraded while every shell it spawned kept
    /// observing. When nash runs, the returned env adds the spool config (NASH.md §6.2)
    /// so the command's events land where the `shell-events` drain looks. Pure given its inputs;
    /// the file-system probes are injectable for tests.
    public static func shellInvocation(
        command: String, nashPath: String,
        environment: [String: String],
        isExecutable: (String) -> Bool = { FileManager.default.isExecutableFile(atPath: $0) },
        policyDisablesNash: () -> Bool = { NashPolicy.load().disable }
    ) -> (path: String, args: [String], environment: [String: String])? {
        if policyDisablesNash() {
            return ("/bin/zsh", ["-c", command], environment)
        }
        guard isExecutable(nashPath) else { return nil }
        var env = environment
        env[ShellSpool.spoolEnvKey] = ShellSpool.defaultDir
        if env[ShellSpool.captureEnvKey] == nil { env[ShellSpool.captureEnvKey] = "tap" }
        return (nashPath, ["-lc", command], env)
    }

    /// The `spawnError` message a nash-less guest replies with instead of running the command under
    /// another shell. Names the fix (rebuild the base) so the failure is actionable, not mysterious.
    public static let nashMissingMessage =
        "nash is not installed in this guest (\(ExecSession.defaultNashPath)) — nash is the required "
        + "agent shell in every virtualized environment (docs/NASH.md §7.5), so the exec surface "
        + "refuses to run without it. Rebuild the macOS VM base image so provisioning installs it."

    // MARK: - The session

    /// Run one exec request to completion: spawn, pump, exit frame. Blocks the calling thread for the
    /// life of the command; the caller owns/closes `fd` afterwards (exec connections are single-use).
    public func run(request: [String: Any]) {
        guard let command = request[AgentWire.Exec.commandKey] as? String else {
            sendFrame([
                AgentWire.Exec.tagKey: AgentWire.Exec.spawnError,
                AgentWire.Exec.messageKey: "exec requires a \"\(AgentWire.Exec.commandKey)\" string",
            ])
            return
        }
        guard
            let shell = Self.shellInvocation(
                command: command, nashPath: nashPath,
                environment: ProcessInfo.processInfo.environment)
        else {
            // nash-only surface: a guest without nash refuses the exec with an actionable error
            // instead of silently running another shell (docs/NASH.md §7.5).
            sendFrame([
                AgentWire.Exec.tagKey: AgentWire.Exec.spawnError,
                AgentWire.Exec.messageKey: Self.nashMissingMessage,
            ])
            return
        }

        var stdinFDs: [Int32] = [-1, -1], stdoutFDs: [Int32] = [-1, -1], stderrFDs: [Int32] = [-1, -1]
        guard pipe(&stdinFDs) == 0, pipe(&stdoutFDs) == 0, pipe(&stderrFDs) == 0 else {
            [stdinFDs, stdoutFDs, stderrFDs].flatMap { $0 }.filter { $0 >= 0 }.forEach { close($0) }
            sendFrame([
                AgentWire.Exec.tagKey: AgentWire.Exec.spawnError,
                AgentWire.Exec.messageKey: "could not allocate stdio pipes (errno \(errno))",
            ])
            return
        }

        let pid = Self.spawnGroupLeader(
            path: shell.path, args: shell.args, environment: shell.environment,
            stdinFD: stdinFDs[0], stdoutFD: stdoutFDs[1], stderrFD: stderrFDs[1])
        // The child ends belong to the child now — closing ours is what makes EOF observable.
        close(stdinFDs[0])
        close(stdoutFDs[1])
        close(stderrFDs[1])
        guard let pid else {
            close(stdinFDs[1])
            close(stdoutFDs[0])
            close(stderrFDs[0])
            sendFrame([
                AgentWire.Exec.tagKey: AgentWire.Exec.spawnError,
                AgentWire.Exec.messageKey: "could not launch the command (posix_spawn errno \(errno))",
            ])
            return
        }

        var stdinFD = stdinFDs[1]
        let outFD = stdoutFDs[0], errFD = stderrFDs[0]
        // Nonblocking stdin so a child that never reads can't park the pump loop mid-write; NOSIGPIPE
        // so a write racing the child's exit gets EPIPE instead of a signal (see init).
        _ = fcntl(stdinFD, F_SETFL, fcntl(stdinFD, F_GETFL, 0) | O_NONBLOCK)
        _ = fcntl(stdinFD, F_SETNOSIGPIPE, 1)

        var stdoutOpen = true, stderrOpen = true
        var exitCode: Int32? = nil
        var exitObservedAt: UInt64? = nil
        var hostHungUp = false
        var scratch = [UInt8](repeating: 0, count: 64 * 1024)

        func closeStdin() {
            if stdinFD >= 0 {
                close(stdinFD)
                stdinFD = -1
            }
        }
        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), stdinFD >= 0
                {
                    Self.writeBounded(fd: stdinFD, data: bytes)
                }
            case AgentWire.Exec.stdinEOF:
                closeStdin()
            case AgentWire.Exec.signal:
                if let n = frame[AgentWire.Exec.signalKey] as? Int {
                    Self.signalGroup(pid: pid, signal: 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) }
        }

        // The pump: one poll over the connection and both output pipes. Ends when the child has been
        // reaped AND both pipes hit EOF — or, if descendants keep the pipes open past the exit, when
        // the post-exit grace lapses (guarantee 1 above).
        pump: while true {
            if exitCode == nil {
                var status: Int32 = 0
                let reaped = waitpid(pid, &status, WNOHANG)
                if reaped == pid {
                    exitCode = Self.exitCode(fromWaitStatus: status)
                    exitObservedAt = Self.nowMillis()
                }
            }
            if exitCode != nil, !stdoutOpen, !stderrOpen { break }
            if let at = exitObservedAt, Self.nowMillis() &- at > UInt64(max(0, graceMillis)) {
                break  // descendants still hold a pipe — the exit frame must not wait for them
            }

            var fds: [pollfd] = [pollfd(fd: fd, events: Int16(POLLIN), revents: 0)]
            if stdoutOpen { fds.append(pollfd(fd: outFD, events: Int16(POLLIN), revents: 0)) }
            if stderrOpen { fds.append(pollfd(fd: errFD, events: Int16(POLLIN), revents: 0)) }
            let ready = poll(&fds, nfds_t(fds.count), 200)
            if ready < 0 {
                if errno == EINTR || errno == EAGAIN { continue }
                Self.signalGroup(pid: pid, signal: SIGKILL)
                break
            }
            for entry in fds where entry.revents != 0 {
                if entry.fd == fd {
                    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 || (n < 0 && errno != EINTR && errno != EAGAIN) {
                        // Host hung up — kill the whole command tree and stop streaming.
                        Self.signalGroup(pid: pid, signal: SIGKILL)
                        hostHungUp = true
                        break pump
                    }
                } else if entry.fd == outFD {
                    stdoutOpen = pumpOutput(from: outFD, tag: AgentWire.Exec.stdout, scratch: &scratch)
                } else if entry.fd == errFD {
                    stderrOpen = pumpOutput(from: errFD, tag: AgentWire.Exec.stderr, scratch: &scratch)
                }
            }
        }

        closeStdin()
        close(outFD)
        close(errFD)
        // If the shell is still running (host hangup path), reap it so the agent leaks no zombies.
        if exitCode == nil {
            var status: Int32 = 0
            if hostHungUp {
                if waitpid(pid, &status, 0) == pid { exitCode = Self.exitCode(fromWaitStatus: status) }
            } else if waitpid(pid, &status, WNOHANG) == pid {
                exitCode = Self.exitCode(fromWaitStatus: status)
            }
        }
        if !hostHungUp {
            sendFrame([
                AgentWire.Exec.tagKey: AgentWire.Exec.exit,
                AgentWire.Exec.codeKey: Int(exitCode ?? -1),
            ])
        }
    }

    /// Drain one readable output pipe into base64 frames. Returns whether the pipe is still open.
    private func pumpOutput(from pipeFD: Int32, tag: String, scratch: inout [UInt8]) -> Bool {
        let n = read(pipeFD, &scratch, scratch.count)
        if n > 0 {
            sendFrame([
                AgentWire.Exec.tagKey: tag,
                AgentWire.Exec.dataKey: Data(bytes: scratch, count: n).base64EncodedString(),
            ])
            return true
        }
        if n < 0, errno == EINTR || errno == EAGAIN { return true }
        return false  // EOF (or unrecoverable error): all writers are gone
    }

    // MARK: - Spawning

    /// `posix_spawn` the shell as its **own process-group leader** with the three pipe fds as its
    /// stdio and everything else closed (`POSIX_SPAWN_CLOEXEC_DEFAULT` — so the child can't inherit
    /// and hold open the vsock connection). Returns the pid, or `nil` with `errno` set.
    static func spawnGroupLeader(
        path: String, args: [String], environment: [String: String],
        stdinFD: Int32, stdoutFD: Int32, stderrFD: Int32
    ) -> pid_t? {
        var attr: posix_spawnattr_t?
        posix_spawnattr_init(&attr)
        defer { posix_spawnattr_destroy(&attr) }
        posix_spawnattr_setflags(
            &attr, Int16(POSIX_SPAWN_SETPGROUP) | Int16(POSIX_SPAWN_CLOEXEC_DEFAULT))
        posix_spawnattr_setpgroup(&attr, 0)  // 0 → the child's own pid becomes its pgid

        var actions: posix_spawn_file_actions_t?
        posix_spawn_file_actions_init(&actions)
        defer { posix_spawn_file_actions_destroy(&actions) }
        posix_spawn_file_actions_adddup2(&actions, stdinFD, 0)
        posix_spawn_file_actions_adddup2(&actions, stdoutFD, 1)
        posix_spawn_file_actions_adddup2(&actions, stderrFD, 2)

        let argv: [UnsafeMutablePointer<CChar>?] =
            ([path] + args).map { strdup($0) } + [nil]
        let envp: [UnsafeMutablePointer<CChar>?] =
            environment.map { strdup("\($0.key)=\($0.value)") } + [nil]
        defer {
            argv.compactMap { $0 }.forEach { free($0) }
            envp.compactMap { $0 }.forEach { free($0) }
        }

        var pid: pid_t = 0
        let rc = posix_spawn(&pid, path, &actions, &attr, argv, envp)
        if rc != 0 {
            errno = rc
            return nil
        }
        return pid
    }

    /// Signal the command's whole process group, falling back to the lone pid (e.g. the group is
    /// already empty of the leader, or SETPGROUP was somehow lost).
    static func signalGroup(pid: pid_t, signal: Int32) {
        if kill(-pid, signal) != 0 {
            _ = kill(pid, signal)
        }
    }

    /// Map a `waitpid` status to the shell convention: exit status, or `128 + signal` for a killed
    /// child (what the host's `mac_vm_exec` reports as the command's exit code).
    static func exitCode(fromWaitStatus status: Int32) -> Int32 {
        if (status & 0x7f) == 0 { return (status >> 8) & 0xff }  // WIFEXITED → WEXITSTATUS
        return 128 + (status & 0x7f)  // WIFSIGNALED → WTERMSIG
    }

    private static func nowMillis() -> UInt64 {
        UInt64(DispatchTime.now().uptimeNanoseconds / 1_000_000)
    }

    // MARK: - Frame I/O

    /// Serialize one frame and write it whole (the pump is single-threaded, so no interleaving).
    private func sendFrame(_ frame: [String: Any]) {
        guard let data = try? JSONSerialization.data(withJSONObject: frame) else { return }
        var line = data
        line.append(0x0A)
        line.withUnsafeBytes { raw in
            var offset = 0
            while offset < raw.count {
                let written = write(fd, raw.baseAddress!.advanced(by: offset), raw.count - offset)
                if written > 0 {
                    offset += written
                } else if written < 0, errno == EINTR {
                    continue
                } else {
                    return  // connection torn down — nothing useful left to do with this frame
                }
            }
        }
    }

    /// Write stdin bytes with a bounded retry: the fd is nonblocking, so a child that stops reading
    /// stalls us at most ~1s before the rest of the payload is dropped (observation must not wedge
    /// the pump; the old FileHandle write could block the whole loop indefinitely).
    static func writeBounded(fd: Int32, data: Data, retries: Int = 500) {
        data.withUnsafeBytes { raw in
            var offset = 0
            var attempts = 0
            while offset < raw.count {
                let written = write(fd, raw.baseAddress!.advanced(by: offset), raw.count - offset)
                if written > 0 {
                    offset += written
                } else if written < 0, errno == EAGAIN || errno == EINTR {
                    attempts += 1
                    if attempts > retries { return }
                    usleep(2000)
                } else {
                    return  // EPIPE etc. — the child is gone; stdin bytes have nowhere to go
                }
            }
        }
    }
}
