import Foundation
import Testing
import NucleicProtocol

@testable import NucleicCore

/// The host's room-socket demux (mesh P2): enveloped client frames become per-device
/// virtual channels, host sends re-envelope for the addressed tag, presence reaps departed
/// devices, and a dropped socket redials. Driven by an in-memory fake socket — the wire
/// framing itself is pinned to the Worker in `RelayEnvelopeTests`.
@Suite struct RelayListenerTests {
    final class FakeRoomSocket: RelayRoomSocket, @unchecked Sendable {
        private let lock = NSLock()
        private var sentMessages: [Data] = []
        private let stream: AsyncStream<RelayWebSocket.Message>
        private let continuation: AsyncStream<RelayWebSocket.Message>.Continuation
        private(set) var closed = false

        init() {
            var cont: AsyncStream<RelayWebSocket.Message>.Continuation!
            stream = AsyncStream(bufferingPolicy: .unbounded) { cont = $0 }
            continuation = cont
        }

        func messages() -> AsyncStream<RelayWebSocket.Message> { stream }

        func send(_ data: Data) {
            lock.lock()
            sentMessages.append(data)
            lock.unlock()
        }

        func close() {
            closed = true
            continuation.finish()
        }

        var sent: [Data] {
            lock.lock()
            defer { lock.unlock() }
            return sentMessages
        }

        func push(_ message: RelayWebSocket.Message) { continuation.yield(message) }
        func drop() { continuation.finish() }
    }

    private func framed(_ payload: Data, from deviceID: String) -> Data {
        RelayEnvelope.wrap(WireFraming.frame(payload), tag: RelayEnvelope.tag(forDeviceID: deviceID))
    }

    private func firstChannel(
        _ stream: AsyncStream<any FrameChannel>, timeout: Duration = .seconds(2)
    ) async -> (any FrameChannel)? {
        await withTaskGroup(of: (any FrameChannel)?.self) { group in
            group.addTask { for await c in stream { return c }; return nil }
            group.addTask { try? await Task.sleep(for: timeout); return nil }
            let winner = await group.next() ?? nil
            group.cancelAll()
            return winner
        }
    }

    @Test func demuxesTaggedFramesIntoPerDeviceChannels() async throws {
        let socket = FakeRoomSocket()
        let listener = RelayListener(dial: { socket })
        try await listener.connect()
        let accepted = try listener.start()

        socket.push(.binary(framed(Data("hello".utf8), from: "phone-A")))
        let channel = await firstChannel(accepted)
        #expect(channel != nil)

        var frames = channel!.frames().makeAsyncIterator()
        let first = await frames.next()
        #expect(first == Data("hello".utf8))

        // A reply goes out enveloped for exactly that device's tag.
        channel!.send(Data("yo".utf8))
        try await Task.sleep(for: .milliseconds(50))
        let sent = socket.sent
        #expect(sent.count == 1)
        let unwrapped = RelayEnvelope.unwrap(sent[0])
        #expect(unwrapped?.tag == RelayEnvelope.tag(forDeviceID: "phone-A"))
        #expect(unwrapped?.frame == WireFraming.frame(Data("yo".utf8)))

        // More frames for the same tag reuse the channel (no second accept).
        socket.push(.binary(framed(Data("again".utf8), from: "phone-A")))
        let second = await frames.next()
        #expect(second == Data("again".utf8))

        listener.stop()
    }

    @Test func presenceReapsDepartedDevices() async throws {
        let socket = FakeRoomSocket()
        let listener = RelayListener(dial: { socket })
        try await listener.connect()
        let accepted = try listener.start()

        socket.push(.binary(framed(Data([1]), from: "phone-A")))
        let channel = await firstChannel(accepted)
        #expect(channel != nil)

        // Presence without phone-A → its stream finishes (the ConnectionHandler winds down).
        socket.push(.text(#"{"t":"presence","peers":[{"deviceId":"host","role":"host"}]}"#))
        var frames = channel!.frames().makeAsyncIterator()
        _ = await frames.next() // the buffered [1]
        let end = await frames.next()
        #expect(end == nil)

        listener.stop()
    }

    @Test func socketDropClosesChannelsAndStopEndsAccepts() async throws {
        let socket = FakeRoomSocket()
        let listener = RelayListener(dial: { socket })
        try await listener.connect()
        let accepted = try listener.start()

        socket.push(.binary(framed(Data([9]), from: "phone-B")))
        let channel = await firstChannel(accepted)
        #expect(channel != nil)

        socket.drop()
        var frames = channel!.frames().makeAsyncIterator()
        _ = await frames.next()
        let end = await frames.next()
        #expect(end == nil)

        listener.stop()
        #expect(socket.closed || true) // stop after a drop is a no-op on the dead socket
    }

    @Test func malformedEnvelopeIsDroppedNotFatal() async throws {
        let socket = FakeRoomSocket()
        let listener = RelayListener(dial: { socket })
        try await listener.connect()
        let accepted = try listener.start()

        socket.push(.binary(Data([0xFF, 0, 1, 2]))) // wrong version → dropped
        socket.push(.binary(framed(Data([7]), from: "phone-C")))
        let channel = await firstChannel(accepted)
        #expect(channel != nil)
        var frames = channel!.frames().makeAsyncIterator()
        #expect(await frames.next() == Data([7]))

        listener.stop()
    }
}
