import Foundation
import Testing

@testable import NucleicCore

/// `JSONRPCConnection`: the JSON-RPC 2.0 transport over a process's stdio (ADAPTERS §0.2), shared
/// by the codex app-server backend (header omitted) and the Grok ACP backend (header present).
/// Driven here against an in-memory `ProcessHandle` — no `codex`/`grok` binary needed.
@Suite struct JSONRPCConnectionTests {
    @Test(.timeLimit(.minutes(1))) func requestResponseCorrelation() async throws {
        let handle = FakeRPCHandle()
        let client = JSONRPCConnection(handle: handle)
        await client.start { _ in }

        async let pending = client.request("initialize", params: ["x": 1])
        // The first request takes id 1; wait until it's been written (and thus registered).
        #expect(await waitUntil { !handle.written().isEmpty })

        // Default dialect matches codex: no "jsonrpc" header, carries id + method + params.
        let outbound = try JSONValue(parsing: handle.written()[0])
        #expect(outbound["jsonrpc"] == nil)
        #expect(outbound["id"]?.intValue == 1)
        #expect(outbound["method"]?.stringValue == "initialize")

        handle.feed(#"{"id":1,"result":{"ok":true}}"#)
        let result = try await pending
        #expect(result["ok"]?.boolValue == true)
    }

    @Test(.timeLimit(.minutes(1))) func includesVersionHeaderWhenRequested() async throws {
        // ACP is standards JSON-RPC 2.0: every outbound message carries "jsonrpc":"2.0".
        let handle = FakeRPCHandle()
        let client = JSONRPCConnection(handle: handle, includeVersionHeader: true)
        await client.start { _ in }

        async let pending = client.request("session/new", params: ["cwd": "/tmp"])
        #expect(await waitUntil { !handle.written().isEmpty })
        let outbound = try JSONValue(parsing: handle.written()[0])
        #expect(outbound["jsonrpc"]?.stringValue == "2.0")
        #expect(outbound["method"]?.stringValue == "session/new")

        handle.feed(#"{"jsonrpc":"2.0","id":1,"result":{"sessionId":"s1"}}"#)
        let result = try await pending
        #expect(result["sessionId"]?.stringValue == "s1")
    }

    @Test(.timeLimit(.minutes(1))) func serverErrorRejectsRequest() async throws {
        let handle = FakeRPCHandle()
        let client = JSONRPCConnection(handle: handle)
        await client.start { _ in }

        async let pending = client.request("turn/start")
        #expect(await waitUntil { !handle.written().isEmpty })
        handle.feed(#"{"id":1,"error":{"code":-32000,"message":"boom"}}"#)

        do {
            _ = try await pending
            Issue.record("expected the server error to reject the request")
        } catch {
            #expect(error is JSONRPCConnection.RPCError)
        }
    }

    @Test(.timeLimit(.minutes(1))) func notificationsArePublished() async throws {
        let handle = FakeRPCHandle()
        let client = JSONRPCConnection(handle: handle)
        await client.start { _ in }

        handle.feed(#"{"method":"turn/started","params":{"threadId":"t1"}}"#)
        var iterator = client.notifications.makeAsyncIterator()
        let note = await iterator.next()
        #expect(note?.method == "turn/started")
        #expect(note?.params["threadId"]?.stringValue == "t1")
    }

    @Test(.timeLimit(.minutes(1))) func endNotificationsDeliversBufferedThenStops() async throws {
        // The ACP drain pattern: session/update notifications are delivered through the read loop,
        // and endNotifications() then terminates the drain loop cleanly (so the backend can settle
        // the turn after the session/prompt response without dropping any update).
        let handle = FakeRPCHandle()
        let client = JSONRPCConnection(handle: handle, includeVersionHeader: true)
        await client.start { _ in }

        handle.feed(#"{"jsonrpc":"2.0","method":"session/update","params":{"n":1}}"#)
        handle.feed(#"{"jsonrpc":"2.0","method":"session/update","params":{"n":2}}"#)

        let collected = IntCollector()
        let drain = Task {
            for await note in client.notifications {
                await collected.append(note.params["n"]?.intValue ?? -1)
            }
        }
        // Both notifications were delivered through the read loop.
        #expect(await waitUntil { await collected.count == 2 })
        // Ending the stream terminates the drain loop.
        await client.endNotifications()
        await drain.value
        #expect(await collected.all == [1, 2])
    }

    @Test(.timeLimit(.minutes(1))) func inboundRequestIsDispatchedAndReplyable() async throws {
        let handle = FakeRPCHandle()
        let client = JSONRPCConnection(handle: handle, includeVersionHeader: true)
        let captured = RequestBox()
        await client.start { request in
            await captured.set(request)
            try? await client.reply(
                to: request.id, result: ["outcome": ["outcome": "selected", "optionId": "allow-once"]])
        }

        handle.feed(
            #"{"jsonrpc":"2.0","id":7,"method":"session/request_permission","params":{"sessionId":"s1"}}"#)
        #expect(await waitUntil { await captured.value != nil })

        let request = await captured.value
        #expect(request?.method == "session/request_permission")
        #expect(request?.params["sessionId"]?.stringValue == "s1")

        // The reply is written back keyed by the same id, with the ACP version header.
        #expect(await waitUntil { handle.written().contains { ($0.replyID) == 7 } })
        let reply = try #require(handle.written().first { $0.replyID == 7 })
        let parsed = try JSONValue(parsing: reply)
        #expect(parsed["jsonrpc"]?.stringValue == "2.0")
        #expect(parsed["result"]?["outcome"]?["optionId"]?.stringValue == "allow-once")
    }

    @Test(.timeLimit(.minutes(1))) func connectionCloseFailsPendingRequests() async throws {
        let handle = FakeRPCHandle()
        let client = JSONRPCConnection(handle: handle)
        await client.start { _ in }

        async let pending = client.request("initialize")
        #expect(await waitUntil { !handle.written().isEmpty })
        handle.finishStdout()  // peer exits without responding

        do {
            _ = try await pending
            Issue.record("expected the closed connection to fail the pending request")
        } catch {
            #expect(error is JSONRPCConnection.RPCError)
        }
    }
}

// MARK: - Test doubles

/// In-memory `ProcessHandle`: feed stdout lines, capture stdin writes.
final class FakeRPCHandle: ProcessHandle, @unchecked Sendable {
    let stdoutLines: AsyncThrowingStream<Data, Error>
    let stderrLines: AsyncThrowingStream<Data, Error>
    private let stdoutContinuation: AsyncThrowingStream<Data, Error>.Continuation
    private let lock = NSLock()
    private var writes: [Data] = []

    var processID: Int32 { 4242 }

    init() {
        var continuation: AsyncThrowingStream<Data, Error>.Continuation!
        stdoutLines = AsyncThrowingStream { continuation = $0 }
        stdoutContinuation = continuation
        stderrLines = AsyncThrowingStream { $0.finish() }
    }

    func feed(_ line: String) { stdoutContinuation.yield(Data(line.utf8)) }
    func finishStdout() { stdoutContinuation.finish() }

    func written() -> [Data] {
        lock.lock(); defer { lock.unlock() }
        return writes
    }

    func writeLine(_ data: Data) throws {
        lock.lock(); defer { lock.unlock() }
        writes.append(data)
    }
    func closeStdin() {}
    func sendSignal(_ sig: Int32) {}
    func wait() async -> Int32 { 0 }
}

/// Holds the most recent inbound request seen by the dispatch handler.
actor RequestBox {
    private(set) var value: JSONRPCConnection.InboundRequest?
    func set(_ request: JSONRPCConnection.InboundRequest) { value = request }
}

/// Accumulates the ints drained from the notification stream (for the endNotifications test).
actor IntCollector {
    private var values: [Int] = []
    func append(_ value: Int) { values.append(value) }
    var count: Int { values.count }
    var all: [Int] { values }
}

extension Data {
    /// The `id` of a JSON-RPC reply message, if this data is one (`{"id":…,"result":…}`).
    fileprivate var replyID: Int? {
        guard let root = try? JSONValue(parsing: self), root["result"] != nil || root["error"] != nil
        else { return nil }
        return root["id"]?.intValue
    }
}

/// Poll `condition` until true or `timeout` elapses (avoids brittle fixed sleeps).
private func waitUntil(
    timeout: Duration = .seconds(2), _ condition: @Sendable () async -> Bool
) async -> Bool {
    let clock = ContinuousClock()
    let deadline = clock.now.advanced(by: timeout)
    while clock.now < deadline {
        if await condition() { return true }
        try? await Task.sleep(for: .milliseconds(5))
    }
    return await condition()
}
