import Foundation
import Testing

@testable import NucleicProtocol

/// The punch state machine (SYNC_PROTOCOL §3.3): deterministic tests feed `Input`s and
/// assert `Output`s — no sockets, no clocks — plus one live test punching two real UDP
/// sockets across loopback (the no-NAT degenerate case both sides' engines must converge on).
@Suite struct PunchEngineTests {
    private let tokenA = Data(repeating: 0xAA, count: PunchWire.tokenLength) // A expects this
    private let tokenB = Data(repeating: 0xBB, count: PunchWire.tokenLength) // B expects this
    private let remote = UDPEndpoint.v4(ip: "203.0.113.7", port: 4_500)

    private func sessionA(candidates: [DirectCandidate]? = nil) -> PunchSession {
        PunchSession(
            localToken: tokenA, remoteToken: tokenB,
            remoteCandidates: candidates ?? [
                DirectCandidate(kind: .reflexive, ip: "203.0.113.7", port: 4_500),
            ])
    }

    @Test func wireRoundTripsAndRejectsJunk() {
        let punch = PunchWire.encode(.punch, token: tokenA)
        #expect(punch.count == PunchWire.datagramLength)
        let decoded = PunchWire.decode(punch)
        #expect(decoded?.kind == .punch)
        #expect(decoded?.token == tokenA)

        #expect(PunchWire.decode(Data([0x00, 0x01])) == nil)
        #expect(PunchWire.decode(Data(repeating: 0, count: PunchWire.datagramLength)) == nil)
        // Right length + magic, unknown type byte.
        var badType = PunchWire.encode(.punch, token: tokenA)
        badType[badType.startIndex + 5] = 0x7F
        #expect(PunchWire.decode(badType) == nil)
    }

    @Test func tickProbesEveryCandidateDeduplicated() {
        var session = sessionA(candidates: [
            DirectCandidate(kind: .ipv6, ip: "2001:db8::1", port: 9_000),
            DirectCandidate(kind: .reflexive, ip: "203.0.113.7", port: 4_500),
            DirectCandidate(kind: .local, ip: "203.0.113.7", port: 4_500), // duplicate endpoint
        ])
        let outputs = session.handle(.tick)
        #expect(outputs.count == 2)
        let expected = PunchWire.encode(.punch, token: tokenB)
        #expect(outputs[0] == .send(expected, to: .v6(ip: "2001:db8::1", port: 9_000)))
        #expect(outputs[1] == .send(expected, to: .v4(ip: "203.0.113.7", port: 4_500)))
    }

    @Test func validPunchAcksToObservedSourceNotCandidate() {
        var session = sessionA()
        // The peer's probe arrives from a NAT-rewritten source no candidate mentioned.
        let observed = UDPEndpoint.v4(ip: "198.51.100.2", port: 61_044)
        let outputs = session.handle(.datagram(PunchWire.encode(.punch, token: tokenA), from: observed))
        #expect(outputs == [
            .send(PunchWire.encode(.ack, token: tokenB), to: observed),
            .inboundValidated(observed),
        ])
        // Repeat punches keep being acked (binding keep-warm) without re-reporting.
        let again = session.handle(.datagram(PunchWire.encode(.punch, token: tokenA), from: observed))
        #expect(again == [.send(PunchWire.encode(.ack, token: tokenB), to: observed)])
    }

    @Test func ackCompletesRoundTripOnce() {
        var session = sessionA()
        let outputs = session.handle(.datagram(PunchWire.encode(.ack, token: tokenA), from: remote))
        #expect(outputs == [.validated(remote)])
        #expect(session.validated == [remote])
        // A duplicate ack (retransmit) validates nothing new.
        #expect(session.handle(.datagram(PunchWire.encode(.ack, token: tokenA), from: remote)).isEmpty)
    }

    @Test func wrongTokenIsIgnoredEntirely() {
        var session = sessionA()
        let wrong = Data(repeating: 0xCC, count: PunchWire.tokenLength)
        #expect(session.handle(.datagram(PunchWire.encode(.punch, token: wrong), from: remote)).isEmpty)
        #expect(session.handle(.datagram(PunchWire.encode(.ack, token: wrong), from: remote)).isEmpty)
        // Sending the peer's token *to us* is also wrong — tokens are direction-bound.
        #expect(session.handle(.datagram(PunchWire.encode(.ack, token: tokenB), from: remote)).isEmpty)
        #expect(session.validated.isEmpty)
    }

    @Test func symmetricBothSidesSkipsOtherwiseAttempts() {
        #expect(!punchWorthAttempting(local: .symmetric, remote: .symmetric))
        #expect(punchWorthAttempting(local: .symmetric, remote: .cone))
        #expect(punchWorthAttempting(local: .cone, remote: .symmetric))
        #expect(punchWorthAttempting(local: .cone, remote: .cone))
        #expect(punchWorthAttempting(local: .unknown, remote: .unknown))
        #expect(punchWorthAttempting(local: .symmetric, remote: .unknown))
    }

    /// Simulated lossy convergence: A's first two probes vanish, the third lands; B's ack
    /// then completes A's round trip. Exactly the retransmit-until-heard behavior the
    /// cadence exists for.
    @Test func lossyPathConvergesViaRetransmit() {
        var a = sessionA()
        var b = PunchSession(
            localToken: tokenB, remoteToken: tokenA,
            remoteCandidates: [DirectCandidate(kind: .reflexive, ip: "198.51.100.1", port: 5_000)])
        let aAddr = UDPEndpoint.v4(ip: "198.51.100.1", port: 5_000)

        // Two ticks lost in transit.
        _ = a.handle(.tick)
        _ = a.handle(.tick)
        // Third tick's probe reaches B.
        guard case .send(let probe, _)? = a.handle(.tick).first else {
            Issue.record("tick produced no probe")
            return
        }
        let bOutputs = b.handle(.datagram(probe, from: aAddr))
        guard case .send(let ack, let ackTo)? = bOutputs.first else {
            Issue.record("valid punch produced no ack")
            return
        }
        #expect(ackTo == aAddr)
        let aOutputs = a.handle(.datagram(ack, from: remote))
        #expect(aOutputs == [.validated(remote)])
    }

    /// Two real engines, two real UDP sockets, loopback: both must prove a round trip.
    /// This is the no-NAT degenerate case, and it exercises the full driver — cadence,
    /// per-socket receive loops, token checks, ack-to-observed-source.
    @Test func liveLoopbackPunchValidatesBothSides() async throws {
        let socketA = try UDPSocket(port: 0, ipv6: false)
        let socketB = try UDPSocket(port: 0, ipv6: false)
        defer {
            socketA.closeSocket()
            socketB.closeSocket()
        }

        let engineA = PunchEngine(
            sockets: .init(v4: socketA, v6: nil),
            localToken: tokenA, remoteToken: tokenB,
            remoteCandidates: [
                DirectCandidate(kind: .local, ip: "127.0.0.1", port: socketB.localPort),
            ])
        let engineB = PunchEngine(
            sockets: .init(v4: socketB, v6: nil),
            localToken: tokenB, remoteToken: tokenA,
            remoteCandidates: [
                DirectCandidate(kind: .local, ip: "127.0.0.1", port: socketA.localPort),
            ])

        async let aFirst: UDPEndpoint? = {
            for await update in await engineA.run(budget: .seconds(5)) where update.roundTrip {
                return update.endpoint
            }
            return nil
        }()
        async let bFirst: UDPEndpoint? = {
            for await update in await engineB.run(budget: .seconds(5)) where update.roundTrip {
                return update.endpoint
            }
            return nil
        }()

        let (aValidated, bValidated) = await (aFirst, bFirst)
        #expect(aValidated == .v4(ip: "127.0.0.1", port: socketB.localPort))
        #expect(bValidated == .v4(ip: "127.0.0.1", port: socketA.localPort))
    }
}
