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
    }

    /// Generations of one device share a routing tag, so a *stale* channel's late close (a
    /// superseded handler's dedup-grace teardown, the liveness reaper) must not evict the
    /// successor channel occupying the slot — that starved the device's live connection and
    /// spawned a duplicate handler on its next frame.
    @Test func staleChannelCloseDoesNotEvictSuccessor() async throws {
        let socket = FakeRoomSocket()
        let listener = RelayListener(dial: { socket })
        try await listener.connect()
        let accepted = try listener.start()

        // Collect every accept on one long-lived consumer (an AsyncStream terminates when a
        // partial iterator is dropped, so per-accept `firstChannel` calls can't see a second one).
        let box = AcceptBox()
        let consumer = Task { for await c in accepted { box.add(c) } }
        defer { consumer.cancel() }

        func waitForAccepts(_ n: Int) async -> Bool {
            for _ in 0..<200 {
                if box.count >= n { return true }
                try? await Task.sleep(for: .milliseconds(10))
            }
            return false
        }

        socket.push(.binary(framed(Data([1]), from: "phone-A")))
        #expect(await waitForAccepts(1))
        let first = box.get(0)

        // The device drops (presence without it) → its slot is reaped…
        socket.push(.text(#"{"t":"presence","peers":[{"deviceId":"host","role":"host"}]}"#))
        // …and reconnects: the next frame opens a fresh channel for the same tag.
        socket.push(.binary(framed(Data([2]), from: "phone-A")))
        #expect(await waitForAccepts(2))
        let second = box.get(1)

        // The stale generation's handler winds down late — this must NOT evict the successor.
        first?.close()
        socket.push(.binary(framed(Data([3]), from: "phone-A")))

        var frames = second!.frames().makeAsyncIterator()
        #expect(await frames.next() == Data([2]))
        #expect(await frames.next() == Data([3]))  // still routed to the live channel
        #expect(box.count == 2)                    // no duplicate third channel was spawned

        listener.stop()
    }

    /// One transient failure at launch must not sideline the relay until restart (the first
    /// connect races LAN/tailnet bring-up by design): connect() throws, but start() keeps
    /// redialing, the listener comes up on a later attempt, and statusChanged reports the
    /// recovery so the Settings row un-freezes.
    @Test func initialConnectFailureRedialsFromStart() async throws {
        struct Flaky: Error {}
        let socket = FakeRoomSocket()
        let attempts = Counter()
        let statuses = StatusBox()
        let listener = RelayListener(
            dial: {
                if attempts.next() == 1 { throw Flaky() }
                return socket
            },
            retryDelay: { _ in .milliseconds(5) })
        listener.statusChanged = { statuses.add($0) }

        await #expect(throws: Flaky.self) { try await listener.connect() }
        #expect(!listener.isConnected)
        let accepted = try listener.start()

        for _ in 0..<200 where !listener.isConnected {
            try await Task.sleep(for: .milliseconds(10))
        }
        #expect(listener.isConnected)

        // The recovered socket is live end-to-end: demux still yields channels and frames.
        socket.push(.binary(framed(Data([4]), from: "phone-D")))
        let channel = await firstChannel(accepted)
        #expect(channel != nil)
        var frames = channel!.frames().makeAsyncIterator()
        #expect(await frames.next() == Data([4]))
        #expect(statuses.last == .some(nil))  // the reconnect was reported

        listener.stop()
    }

    /// The field failure this guards: the redial loop dies without a socket (here: its task
    /// gets cancelled mid-backoff) and the listener would sit "Retrying…" forever with no
    /// dial traffic. The watchdog must notice the dead loop and spawn a fresh one.
    @Test func watchdogRevivesADeadRedialLoop() async throws {
        let socket = FakeRoomSocket()
        let attempts = Counter()
        let listener = RelayListener(
            dial: { socket },
            retryDelay: { _ in
                // First loop: kill the redial task from inside its own backoff — the loop
                // exits via its cancellation check without ever dialing, exactly like a
                // leaked/died loop in production.
                if attempts.next() == 1 {
                    withUnsafeCurrentTask { $0?.cancel() }
                    return .seconds(60)
                }
                return .milliseconds(1)
            },
            watchdogInterval: .milliseconds(20))
        _ = try listener.start()  // no connect() — starts disconnected, redialing

        for _ in 0..<200 where !listener.isConnected {
            try await Task.sleep(for: .milliseconds(10))
        }
        #expect(listener.isConnected)

        listener.stop()
    }

    /// Redial is single-flight: the watchdog firing while a (slow) redial loop is already
    /// live must not stack a second loop — concurrent loops double-dial and double-report.
    @Test func watchdogDoesNotStackRedialLoops() async throws {
        struct Down: Error {}
        let attempts = Counter()
        let listener = RelayListener(
            dial: {
                _ = attempts.next()
                try? await Task.sleep(for: .milliseconds(40))
                throw Down()
            },
            retryDelay: { _ in .milliseconds(40) },
            watchdogInterval: .milliseconds(10))
        _ = try listener.start()

        // ~4 watchdog periods per dial cycle: were loops stacking, attempts would multiply
        // well past the single-loop cadence (~1 per 80ms).
        try await Task.sleep(for: .milliseconds(400))
        listener.stop()
        #expect(attempts.total <= 7)
    }

    /// Coming back online must redial the room socket *now*, not sit out the redial backoff
    /// (up to 30s) or the watchdog interval — that idle window is exactly when a cellular/WAN
    /// device finds the room hostless. The loop is parked on a 60s backoff; `networkPathChanged`
    /// supersedes it with an immediate dial that connects well inside that window.
    @Test func networkPathChangeRedialsImmediately() async throws {
        let socket = FakeRoomSocket()
        let attempts = Counter()
        let listener = RelayListener(
            dial: { _ = attempts.next(); return socket },
            retryDelay: { _ in .seconds(60) },   // a parked loop can't dial in the test window
            watchdogInterval: .seconds(60))       // rule out the watchdog doing the reconnect
        _ = try listener.start()                  // no connect() → parked in backoff, disconnected

        // Let the redial loop reach its 60s sleep, and confirm it hasn't dialed on its own.
        try await Task.sleep(for: .milliseconds(50))
        #expect(!listener.isConnected)
        #expect(attempts.total == 0)

        listener.networkPathChanged()
        for _ in 0..<200 where !listener.isConnected {
            try await Task.sleep(for: .milliseconds(10))
        }
        #expect(listener.isConnected)

        // A kick while already connected is a no-op — it must not recycle the live socket.
        listener.networkPathChanged()
        try await Task.sleep(for: .milliseconds(30))
        #expect(listener.isConnected)
        #expect(!socket.closed)

        listener.stop()
    }

    final class Counter: @unchecked Sendable {
        private let lock = NSLock()
        private var value = 0
        func next() -> Int {
            lock.lock(); defer { lock.unlock() }
            value += 1
            return value
        }
        var total: Int { lock.lock(); defer { lock.unlock() }; return value }
    }

    final class StatusBox: @unchecked Sendable {
        private let lock = NSLock()
        private var statuses: [String?] = []
        func add(_ s: String?) { lock.lock(); statuses.append(s); lock.unlock() }
        var last: String?? { lock.lock(); defer { lock.unlock() }; return statuses.last }
    }

    final class AcceptBox: @unchecked Sendable {
        private let lock = NSLock()
        private var channels: [any FrameChannel] = []
        func add(_ c: any FrameChannel) { lock.lock(); channels.append(c); lock.unlock() }
        func get(_ i: Int) -> (any FrameChannel)? {
            lock.lock(); defer { lock.unlock() }
            return channels.count > i ? channels[i] : nil
        }
        var count: Int { lock.lock(); defer { lock.unlock() }; return channels.count }
    }

    @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()
    }
}
