import Foundation
import Testing
#if canImport(CryptoKit)
import CryptoKit
#else
import Crypto
#endif

@testable import NucleicProtocol

/// The multipath mux (SYNC_PROTOCOL §3.3) — the correctness-critical piece of the direct
/// upgrade. These tests run two muxes over in-memory pipes and assert the ordering
/// invariants the Noise nonce discipline depends on: passthrough-until-armed, in-order
/// migration at the SWITCH marker, *lossless* failback (in-flight direct frames black-holed
/// on the wire must reappear via relay retransmit, exactly once, in order), forged-control
/// rejection, and re-upgrade after a failback.
@Suite struct MultipathChannelTests {
    // MARK: In-memory pipe

    /// An ordered, paired FrameChannel like the CoreTests `MemoryChannel`, plus a
    /// `blackhole()` switch that models in-flight loss: frames sent to this end after the
    /// call vanish (the sender doesn't know), which is precisely the failure the RESYNC
    /// retransmit exists to repair.
    final class PipeChannel: FrameChannel, @unchecked Sendable {
        let name: String
        private let hint: SyncTransportHint?
        private var peer: PipeChannel?
        private let lock = NSLock()
        private var continuation: AsyncStream<Data>.Continuation?
        private let stream: AsyncStream<Data>
        private var dropInbound = false
        private var finished = false

        init(name: String, transport: SyncTransportHint?) {
            self.name = name
            self.hint = transport
            var continuation: AsyncStream<Data>.Continuation?
            self.stream = AsyncStream { continuation = $0 }
            self.continuation = continuation
        }

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

        var peerDescription: String { name }
        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() }
            guard !dropInbound, !finished else { return }
            continuation?.yield(frame)
        }

        /// Everything delivered to this end from now on is lost in flight.
        func blackhole() {
            lock.lock()
            defer { lock.unlock() }
            dropInbound = true
        }

        func close() {
            finish()
            peer?.finish()
        }

        private func finish() {
            lock.lock()
            defer { lock.unlock() }
            guard !finished else { return }
            finished = true
            continuation?.finish()
        }
    }

    actor FrameLog {
        private(set) var frames: [Data] = []
        func append(_ frame: Data) { frames.append(frame) }
    }

    private func collect(_ mux: MultipathFrameChannel) -> (log: FrameLog, task: Task<Void, Never>) {
        let log = FrameLog()
        let task = Task {
            for await frame in mux.frames() {
                await log.append(frame)
            }
        }
        return (log, task)
    }

    private func waitFor(
        _ log: FrameLog, count: Int, timeout: Duration = .seconds(5)
    ) async -> [Data] {
        let clock = ContinuousClock()
        let deadline = clock.now + timeout
        while clock.now < deadline {
            let frames = await log.frames
            if frames.count >= count { return frames }
            try? await clock.sleep(for: .milliseconds(10))
        }
        return await log.frames
    }

    private func frame(_ i: Int) -> Data {
        Data("frame-\(i)".utf8)
    }

    private let key = MultipathFrameChannel.bindingKey(handshakeHash: Data(repeating: 1, count: 32))
    private let wrongKey = MultipathFrameChannel.bindingKey(handshakeHash: Data(repeating: 2, count: 32))

    /// A fully upgraded A↔B mux pair over relay + direct pipes.
    private func upgradedPair() async throws -> (
        a: MultipathFrameChannel, b: MultipathFrameChannel,
        bLog: FrameLog, aLog: FrameLog,
        directA: PipeChannel, directB: PipeChannel
    ) {
        let (relayA, relayB) = PipeChannel.pair(transport: .relay)
        let a = MultipathFrameChannel(primary: relayA)
        let b = MultipathFrameChannel(primary: relayB)
        a.arm(bindingKey: key)
        b.arm(bindingKey: key)
        let (aC, _) = collect(a) // drain a's inbound too
        _ = aC
        let (bLog, _) = collect(b)
        let aLog = aC

        let (directA, directB) = PipeChannel.pair(transport: .direct)
        a.attachDirect(directA)
        b.attachDirect(directB)
        async let aSwitched = a.verifyAndSwitch(timeout: .seconds(3))
        async let bSwitched = b.verifyAndSwitch(timeout: .seconds(3))
        let (sa, sb) = await (aSwitched, bSwitched)
        try #require(sa && sb)
        return (a, b, bLog, aLog, directA, directB)
    }

    // MARK: Tests

    @Test func unarmedMuxIsPurePassthrough() async {
        let (relayA, relayB) = PipeChannel.pair(transport: .relay)
        let mux = MultipathFrameChannel(primary: relayA)
        let (log, _) = collect(mux)
        let (peerLog, peerTask) = collect(MultipathFrameChannel(primary: relayB))
        _ = peerTask

        mux.send(frame(1))
        // Even a frame shaped exactly like a control frame passes through unarmed —
        // interpretation requires the arm() key.
        let controlLookalike = MultipathFrameChannel.Control.encode(
            .switchToDirect, generation: 1, value: 1, key: key)
        relayB.send(controlLookalike)
        relayB.send(frame(2))

        let received = await waitFor(log, count: 2)
        #expect(received == [controlLookalike, frame(2)])
        let peerReceived = await waitFor(peerLog, count: 1)
        #expect(peerReceived == [frame(1)])
        #expect(mux.transport == .relay)
        #expect(mux.activePath == .relay)
    }

    @Test func switchMigratesWithoutReorderingAcrossThePaths() async throws {
        let (relayA, relayB) = PipeChannel.pair(transport: .relay)
        let a = MultipathFrameChannel(primary: relayA)
        let b = MultipathFrameChannel(primary: relayB)
        a.arm(bindingKey: key)
        b.arm(bindingKey: key)
        let (bLog, _) = collect(b)
        let (_, _) = collect(a)

        // Traffic flowing before, during, and after the upgrade must arrive in order.
        a.send(frame(1))
        a.send(frame(2))

        let (directA, directB) = PipeChannel.pair(transport: .direct)
        a.attachDirect(directA)
        b.attachDirect(directB)
        async let aSwitched = a.verifyAndSwitch(timeout: .seconds(3))
        async let bSwitched = b.verifyAndSwitch(timeout: .seconds(3))
        let (sa, sb) = await (aSwitched, bSwitched)
        try #require(sa && sb)
        #expect(a.activePath == .direct)

        a.send(frame(3))
        a.send(frame(4))

        let received = await waitFor(bLog, count: 4)
        #expect(received == [frame(1), frame(2), frame(3), frame(4)])
    }

    @Test func inFlightLossIsRepairedByResyncRetransmit() async throws {
        let (a, _, bLog, _, _, directB) = try await upgradedPair()

        a.send(frame(1))
        a.send(frame(2))
        _ = await waitFor(bLog, count: 2)

        // Frames 3 and 4 vanish on the direct wire; then the path dies. B has processed
        // through 2; A's retransmit buffer still holds 3 and 4.
        directB.blackhole()
        a.send(frame(3))
        a.send(frame(4))
        directB.close()

        // Both sides detect, RESYNC over relay, and A retransmits exactly the gap.
        let received = await waitFor(bLog, count: 4)
        #expect(received == [frame(1), frame(2), frame(3), frame(4)])
        #expect(a.activePath == .relay)

        // The session is still live on relay: post-failback traffic flows.
        a.send(frame(5))
        let after = await waitFor(bLog, count: 5)
        #expect(after.last == frame(5))
    }

    @Test func sendsDuringFailbackAreQueuedNotLost() async throws {
        let (a, _, bLog, _, directA, directB) = try await upgradedPair()

        a.send(frame(1))
        _ = await waitFor(bLog, count: 1)

        // Kill the path and immediately keep sending — these race the RESYNC exchange
        // and must come out after the retransmitted gap, in order.
        directB.blackhole()
        a.send(frame(2))
        directA.close()
        a.send(frame(3))
        a.send(frame(4))

        let received = await waitFor(bLog, count: 4)
        #expect(received == [frame(1), frame(2), frame(3), frame(4)])
    }

    @Test func forgedControlFramesAreDroppedNotObeyed() async throws {
        let (relayA, relayB) = PipeChannel.pair(transport: .relay)
        let a = MultipathFrameChannel(primary: relayA)
        a.arm(bindingKey: key)
        let (log, _) = collect(a)

        // A "relay" forging control frames without the Noise-derived key: dropped, silently.
        relayB.send(MultipathFrameChannel.Control.encode(
            .switchToDirect, generation: 1, value: 1, key: wrongKey))
        relayB.send(MultipathFrameChannel.Control.encode(
            .resync, generation: 1, value: 0, key: wrongKey))
        relayB.send(frame(1))

        let received = await waitFor(log, count: 1)
        #expect(received == [frame(1)]) // forgeries neither surfaced nor obeyed
        #expect(a.activePath == .relay)
    }

    @Test func reUpgradeAfterFailbackContinuesInOrder() async throws {
        let (a, b, bLog, _, _, directB) = try await upgradedPair()

        a.send(frame(1))
        _ = await waitFor(bLog, count: 1)
        directB.close() // clean path death, nothing in flight
        _ = await waitFor(bLog, count: 1)

        // Wait for both sides to settle back on relay before re-upgrading.
        let clock = ContinuousClock()
        let deadline = clock.now + .seconds(5)
        while (a.activePath != .relay || !a.canAttachDirect || !b.canAttachDirect),
              clock.now < deadline
        {
            try? await clock.sleep(for: .milliseconds(20))
        }
        try #require(a.canAttachDirect && b.canAttachDirect)

        a.send(frame(2)) // rides relay between generations

        let (directA2, directB2) = PipeChannel.pair(transport: .direct)
        a.attachDirect(directA2)
        b.attachDirect(directB2)
        async let aSwitched = a.verifyAndSwitch(timeout: .seconds(3))
        async let bSwitched = b.verifyAndSwitch(timeout: .seconds(3))
        let (sa, sb) = await (aSwitched, bSwitched)
        try #require(sa && sb)

        a.send(frame(3))
        let received = await waitFor(bLog, count: 3)
        #expect(received == [frame(1), frame(2), frame(3)])
        #expect(a.activePath == .direct)
    }

    @Test func relayLegDeathTearsDownTheWholeSession() async throws {
        let (a, _, bLog, _, _, _) = try await upgradedPair()
        _ = bLog

        let finishTask = Task { () -> Bool in
            for await _ in a.frames() {}
            return true
        }
        // The spine dies: the mux must finish its stream even though direct is live.
        a.close()
        let sawFinish = await finishTask.value
        #expect(sawFinish)
    }

    @Test func verifyTimeoutDetachesAndStaysOnRelay() async throws {
        let (relayA, relayB) = PipeChannel.pair(transport: .relay)
        let a = MultipathFrameChannel(primary: relayA)
        a.arm(bindingKey: key)
        let (bLog, _) = collect(MultipathFrameChannel(primary: relayB))

        // Attach a direct leg whose peer never answers probes (nothing armed over there).
        let (directA, _) = PipeChannel.pair(transport: .direct)
        a.attachDirect(directA)
        let switched = await a.verifyAndSwitch(timeout: .milliseconds(200))
        #expect(!switched)
        #expect(a.activePath == .relay)

        // Frames still flow on relay, and a fresh attach is allowed again.
        a.send(frame(1))
        let received = await waitFor(bLog, count: 1)
        #expect(received == [frame(1)])
        #expect(a.canAttachDirect)
    }
}
