import Darwin
import Foundation
import Testing

@testable import VMAgentCore

/// Integration tests for ``ExecSession`` over a real `socketpair` — the exact code path the vsock
/// exec connection runs, minus the vsock. Each test plays the host: writes the request line (plus any
/// pipelined frames), then reads NDJSON frames off its end until the terminal `exit`/`spawnError`
/// frame or a deadline.
///
/// The wedge cases (`backgroundChild…`, `signalKills…`) are the regression suite for the
/// "mac_vm_exec never returns" hang (docs/NASH.md §7.5): under the old Foundation-`Process` loop each
/// of them held the exit frame off until every pipe-holding descendant died.
/// `.serialized`: each test spawns real processes, signals process groups, and juggles raw fds —
/// process-global state that parallel execution can only make flakier, for no coverage gain.
@Suite(.serialized) struct ExecSessionTests {

    /// One driven exec: the frames received (in order) and whether the terminal frame arrived in time.
    private struct Outcome {
        var frames: [[String: Any]] = []
        var sawTerminal = false

        var exitCode: Int? {
            for frame in frames where frame[AgentWire.Exec.tagKey] as? String == AgentWire.Exec.exit {
                return frame[AgentWire.Exec.codeKey] as? Int
            }
            return nil
        }

        func output(_ tag: String) -> String {
            var data = Data()
            for frame in frames where frame[AgentWire.Exec.tagKey] as? String == tag {
                if let b64 = frame[AgentWire.Exec.dataKey] as? String,
                    let bytes = Data(base64Encoded: b64)
                {
                    data.append(bytes)
                }
            }
            return String(decoding: data, as: UTF8.self)
        }
    }

    /// Run one session against `request`, optionally streaming extra host frames after a short delay,
    /// and collect frames until the terminal one (or `deadline` seconds).
    ///
    /// `nashPath` defaults to `/bin/zsh` as a stand-in nash: the exec surface is nash-only (a missing
    /// binary is a refusal, not a fallback — see `nashLessGuestRefusesTheExec`), and the session just
    /// spawns `<nashPath> -lc <command>`, which zsh accepts — so the spawn/pump/signal mechanics under
    /// test run against a shell every test machine has, independent of a locally installed nash.
    private func drive(
        request: [String: Any], laterFrames: [[String: Any]] = [],
        laterDelayMillis: UInt32 = 300, deadline: TimeInterval = 15,
        nashPath: String = "/bin/zsh", pipeGraceMillis: Int32 = 1500
    ) throws -> Outcome {
        var pair: [Int32] = [0, 0]
        try #require(socketpair(AF_UNIX, SOCK_STREAM, 0, &pair) == 0)
        let sessionFD = pair[0], hostFD = pair[1]

        // The session thread must be JOINED before this helper returns: tests run in parallel, so a
        // straggling `close(sessionFD)` after return could close a reused fd belonging to another
        // in-flight test (fd-reuse corruption presenting as a flaky crashed test process).
        let sessionDone = DispatchSemaphore(value: 0)
        let thread = Thread {
            ExecSession(fd: sessionFD, nashPath: nashPath, pipeGraceMillis: pipeGraceMillis)
                .run(request: request)
            close(sessionFD)
            sessionDone.signal()
        }
        thread.start()

        func writeLine(_ object: [String: Any]) {
            guard var data = try? JSONSerialization.data(withJSONObject: object) else { return }
            data.append(0x0A)
            data.withUnsafeBytes { _ = write(hostFD, $0.baseAddress, $0.count) }
        }
        writeLine(request)
        if !laterFrames.isEmpty {
            usleep(laterDelayMillis * 1000)
            for frame in laterFrames { writeLine(frame) }
        }

        var outcome = Outcome()
        var buffer = LineSplitBuffer()
        var scratch = [UInt8](repeating: 0, count: 64 * 1024)
        let end = Date().addingTimeInterval(deadline)
        while Date() < end, !outcome.sawTerminal {
            var pfd = pollfd(fd: hostFD, events: Int16(POLLIN), revents: 0)
            let ready = poll(&pfd, 1, 200)
            guard ready > 0, (pfd.revents & Int16(POLLIN)) != 0 else { continue }
            let n = read(hostFD, &scratch, scratch.count)
            if n <= 0 { break }
            buffer.append(Data(bytes: scratch, count: n))
            while let line = buffer.nextLine() {
                guard !line.isEmpty,
                    let obj = (try? JSONSerialization.jsonObject(with: line)) as? [String: Any]
                else { continue }
                outcome.frames.append(obj)
                let tag = obj[AgentWire.Exec.tagKey] as? String
                if tag == AgentWire.Exec.exit || tag == AgentWire.Exec.spawnError {
                    outcome.sawTerminal = true
                }
            }
        }
        // Hanging up unblocks a session still draining (its host-hangup path kills the command
        // group), then the join above's counterpart: wait for the session thread to finish so its
        // fd close cannot race a parallel test.
        close(hostFD)
        _ = sessionDone.wait(timeout: .now() + 20)
        return outcome
    }

    private func execRequest(_ command: String) -> [String: Any] {
        ["op": AgentWire.Exec.op, AgentWire.Exec.commandKey: command]
    }

    // MARK: - Baseline

    @Test func simpleCommandStreamsOutputAndExit() throws {
        let outcome = try drive(request: execRequest("echo out-here; echo err-here 1>&2; exit 3"))
        #expect(outcome.sawTerminal)
        #expect(outcome.exitCode == 3)
        #expect(outcome.output(AgentWire.Exec.stdout).contains("out-here"))
        #expect(outcome.output(AgentWire.Exec.stderr).contains("err-here"))
    }

    @Test func malformedRequestGetsSpawnError() throws {
        let outcome = try drive(request: ["op": AgentWire.Exec.op])  // no command
        #expect(outcome.sawTerminal)
        let tags = outcome.frames.compactMap { $0[AgentWire.Exec.tagKey] as? String }
        #expect(tags.contains(AgentWire.Exec.spawnError))
    }

    @Test func stdinFramesReachTheChild() throws {
        let outcome = try drive(
            request: execRequest("cat"),
            laterFrames: [
                [
                    AgentWire.Exec.tagKey: AgentWire.Exec.stdin,
                    AgentWire.Exec.dataKey: Data("ping-pong\n".utf8).base64EncodedString(),
                ],
                [AgentWire.Exec.tagKey: AgentWire.Exec.stdinEOF],
            ])
        #expect(outcome.sawTerminal)
        #expect(outcome.exitCode == 0)
        #expect(outcome.output(AgentWire.Exec.stdout).contains("ping-pong"))
    }

    // MARK: - The wedge cases (docs/NASH.md §7.5 — "mac_vm_exec never returns")

    /// A backgrounded child inherits the stdout write-end and outlives the shell. The exit frame must
    /// still arrive within the pipe grace — under the old implementation it waited the full 30s.
    @Test func backgroundChildCannotHoldTheExitFrame() throws {
        let started = Date()
        let outcome = try drive(
            request: execRequest("sleep 30 & echo started"),
            pipeGraceMillis: 1000)
        #expect(outcome.sawTerminal)
        #expect(outcome.exitCode == 0)
        #expect(outcome.output(AgentWire.Exec.stdout).contains("started"))
        #expect(Date().timeIntervalSince(started) < 10)  // grace (1s) + slack, nowhere near 30s
    }

    /// A `signal` frame must kill the whole process group — including pipeline children the shell
    /// forked — and produce a prompt exit frame, not an orphan holding the pipes for 30s.
    @Test func signalFrameKillsTheWholeProcessGroup() throws {
        let started = Date()
        let outcome = try drive(
            request: execRequest("sleep 30 | sleep 31"),
            laterFrames: [[AgentWire.Exec.tagKey: AgentWire.Exec.signal, AgentWire.Exec.signalKey: 9]])
        #expect(outcome.sawTerminal)
        #expect(outcome.exitCode == 128 + 9)
        #expect(Date().timeIntervalSince(started) < 10)
    }

    // MARK: - Spawn mechanics

    @Test func spawnedShellIsItsOwnProcessGroupLeader() throws {
        var fds: [Int32] = [-1, -1]
        try #require(pipe(&fds) == 0)
        let pid = ExecSession.spawnGroupLeader(
            path: "/bin/sleep", args: ["5"], environment: [:],
            stdinFD: fds[0], stdoutFD: fds[1], stderrFD: fds[1])
        close(fds[0])
        close(fds[1])
        let spawned = try #require(pid)
        #expect(getpgid(spawned) == spawned)  // SETPGROUP took: the child leads its own group
        ExecSession.signalGroup(pid: spawned, signal: SIGKILL)
        var status: Int32 = 0
        #expect(waitpid(spawned, &status, 0) == spawned)
        #expect(ExecSession.exitCode(fromWaitStatus: status) == 128 + SIGKILL)
    }

    @Test func exitCodeMapsWaitStatusConventionally() {
        #expect(ExecSession.exitCode(fromWaitStatus: 0) == 0)
        #expect(ExecSession.exitCode(fromWaitStatus: 3 << 8) == 3)  // WIFEXITED status 3
        #expect(ExecSession.exitCode(fromWaitStatus: SIGKILL) == 128 + SIGKILL)  // signaled
    }

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

    /// A guest without nash refuses the exec outright — terminal `spawnError` carrying the
    /// actionable rebuild-the-base message, no zsh fallback, no hang.
    @Test func nashLessGuestRefusesTheExec() throws {
        let outcome = try drive(
            request: execRequest("echo hi"), nashPath: "/nonexistent/nash")
        #expect(outcome.sawTerminal)
        #expect(outcome.exitCode == nil)
        let error = outcome.frames.first {
            $0[AgentWire.Exec.tagKey] as? String == AgentWire.Exec.spawnError
        }
        #expect((error?[AgentWire.Exec.messageKey] as? String) == ExecSession.nashMissingMessage)
    }

    @Test func shellInvocationIsNashOnly() {
        let nash = "/usr/local/bin/nash"
        // nash installed → nash -lc with the spool env wired in. No request opt-in exists: nash is
        // the only shell the exec surface runs (docs/NASH.md §7.5 — locked).
        let picked = ExecSession.shellInvocation(
            command: "echo hi", nashPath: nash, environment: ["A": "1"],
            isExecutable: { $0 == nash }, policyDisablesNash: { false })
        #expect(picked?.path == nash)
        #expect(picked?.args == ["-lc", "echo hi"])
        #expect(picked?.environment[ShellSpool.spoolEnvKey] == ShellSpool.defaultDir)
        #expect(picked?.environment[ShellSpool.captureEnvKey] == "tap")

        // No nash → the exec is REFUSED (nil → spawnError), never silently degraded to zsh.
        let absent = ExecSession.shellInvocation(
            command: "echo hi", nashPath: nash, environment: ["A": "1"],
            isExecutable: { _ in false }, policyDisablesNash: { false })
        #expect(absent == nil)
        #expect(ExecSession.nashMissingMessage.contains(nash))

        // The in-guest operator break-glass is the sole exception: the trusted policy's kill
        // switch → zsh. It reads the root-owned file, NOT the environment — an agent that exports
        // `NUCLEIC_NASH_DISABLE=1` gets nash anyway (docs/NASH.md §4.3).
        let disabled = ExecSession.shellInvocation(
            command: "x", nashPath: nash, environment: [:], isExecutable: { _ in true },
            policyDisablesNash: { true })
        #expect(disabled?.path == "/bin/zsh")
        #expect(disabled?.args == ["-c", "x"])

        let envLever = ExecSession.shellInvocation(
            command: "x", nashPath: nash,
            environment: ["NUCLEIC_NASH_DISABLE": "1"], isExecutable: { _ in true },
            policyDisablesNash: { false })
        #expect(envLever?.path == nash)

        // An existing capture preference is preserved, not clobbered to tap.
        let meta = ExecSession.shellInvocation(
            command: "x", nashPath: nash,
            environment: [ShellSpool.captureEnvKey: "meta"], isExecutable: { _ in true },
            policyDisablesNash: { false })
        #expect(meta?.environment[ShellSpool.captureEnvKey] == "meta")
    }
}
