import Foundation
import Testing

@testable import NucleicProtocol

/// The full Covalence direct upgrade end to end on loopback (SYNC §3.3): the client driver
/// and host upgrader exchange offer → answer over an in-memory "relay", both punch real UDP
/// sockets, bring up RUDP over the punched sockets, and migrate a live mux session onto the
/// direct path — asserting frames sent before, during, and after the switch arrive in order
/// with no loss. This replaces the QUIC go/no-go gate: the RUDP path is pure POSIX, so the
/// whole thing runs anywhere (Linux CI included).
@Suite(.serialized) struct DirectUpgradeLoopbackTests {
    /// In-memory ordered relay channel pair (mirrors the CoreTests `MemoryChannel`).
    final class RelayPipe: FrameChannel, @unchecked Sendable {
        private let hint: SyncTransportHint?
        private var peer: RelayPipe?
        private let lock = NSLock()
        private var cont: AsyncStream<Data>.Continuation?
        private let stream: AsyncStream<Data>
        private var finished = false

        init(transport: SyncTransportHint?) {
            self.hint = transport
            var c: AsyncStream<Data>.Continuation?
            self.stream = AsyncStream { c = $0 }
            self.cont = c
        }

        static func pair(transport: SyncTransportHint?) -> (RelayPipe, RelayPipe) {
            let a = RelayPipe(transport: transport)
            let b = RelayPipe(transport: transport)
            a.peer = b
            b.peer = a
            return (a, b)
        }

        var peerDescription: String { "relay" }
        var transport: SyncTransportHint? { hint }
        func frames() -> AsyncStream<Data> { stream }
        func send(_ frame: Data) { peer?.deliver(frame) }
        private func deliver(_ frame: Data) {
            lock.lock(); defer { lock.unlock() }
            if !finished { cont?.yield(frame) }
        }
        func close() {
            finish(); peer?.finish()
        }
        private func finish() {
            lock.lock(); defer { lock.unlock() }
            if !finished { finished = true; cont?.finish() }
        }
    }

    @Test func upgradeMigratesLiveSessionInOrder() async throws {
        // Point candidate gathering at an in-process STUN responder (no public servers).
        let stun = try StunResponder()
        defer { stun.stop() }
        setenv("NUCLEIC_STUN_SERVERS", "127.0.0.1:\(stun.port)", 1)
        defer { unsetenv("NUCLEIC_STUN_SERVERS") }

        let (relayClient, relayHost) = RelayPipe.pair(transport: .relay)
        let clientMux = MultipathFrameChannel(primary: relayClient)
        let hostMux = MultipathFrameChannel(primary: relayHost)
        let key = MultipathFrameChannel.bindingKey(handshakeHash: Data(repeating: 5, count: 32))
        clientMux.arm(bindingKey: key)
        hostMux.arm(bindingKey: key)

        let hostSeen = FrameSink()
        let collect = Task { for await f in hostMux.frames() { await hostSeen.add(f) } }
        let clientDrain = Task { for await _ in clientMux.frames() {} }
        defer {
            collect.cancel()
            clientDrain.cancel()
        }

        let replies = DirectReplyBox()
        let host = DirectHostUpgrader(
            mux: hostMux,
            sendAnswer: { replies.deliver(.answer($0)) },
            sendGo: { replies.deliver(.go) },
            sendDecline: { replies.deliver(.decline($0)) })

        // Traffic before the upgrade.
        clientMux.send(Data("before-1".utf8))
        clientMux.send(Data("before-2".utf8))

        let upgraded = await DirectUpgradeDriver.run(
            mux: clientMux, attempts: 1, retryDelay: .seconds(1), replyTimeout: .seconds(15),
            replies: replies
        ) { msg in
            // Detached like the real ConnectionHandler → SyncHost path: the host's punch /
            // verify run concurrently with the client's, rather than blocking the driver.
            switch msg {
            case .directOffer(let offer): Task { await host.start(remoteOffer: offer) }
            case .directSelect(let sel): Task { await host.select(sel) }
            default: break
            }
        }
        try #require(upgraded, "loopback direct upgrade must complete")
        #expect(clientMux.activePath == .direct)

        // Traffic after the switch — must land after the pre-switch frames, in order.
        clientMux.send(Data("after-1".utf8))
        clientMux.send(Data("after-2".utf8))

        let frames = await hostSeen.wait(count: 4, timeout: .seconds(10))
        #expect(frames == [
            Data("before-1".utf8), Data("before-2".utf8),
            Data("after-1".utf8), Data("after-2".utf8),
        ])

        await host.cancel()
        clientMux.close()
        hostMux.close()
    }

    /// A minimal in-process STUN responder (reflects the observed source as XOR-MAPPED-
    /// ADDRESS) so candidate gathering classifies against loopback.
    final class StunResponder: @unchecked Sendable {
        let port: UInt16
        private let socket: UDPSocket
        private let task: Task<Void, Never>

        init() throws {
            let s = try UDPSocket(port: 0, ipv6: false)
            self.socket = s
            self.port = s.localPort
            self.task = Task {
                let clock = ContinuousClock()
                while !Task.isCancelled {
                    guard let packet = try? await s.receive(deadline: clock.now + .seconds(30))
                    else { continue }
                    guard packet.data.count >= 20 else { continue }
                    let txid = Data(packet.data.suffix(12))
                    var r = Data([0x01, 0x01, 0x00, 0x0C, 0x21, 0x12, 0xA4, 0x42])
                    r += txid
                    r += Data([0x00, 0x20, 0x00, 0x08])
                    let xport = packet.from.port ^ UInt16(StunMessage.magicCookie >> 16)
                    r += Data([0x00, 0x01, UInt8(xport >> 8), UInt8(xport & 0xFF)])
                    let octets = packet.from.ip.split(separator: ".").compactMap { UInt8($0) }
                    r += Data(zip(octets, [0x21, 0x12, 0xA4, 0x42]).map(^))
                    try? s.send(r, to: packet.from)
                }
            }
        }

        func stop() {
            task.cancel()
            socket.closeSocket()
        }
    }
}
