import Foundation
import Testing

@testable import NucleicProtocol

/// The STUN codec (SYNC_PROTOCOL §3.3): request shape and (XOR-)MAPPED-ADDRESS decoding are
/// pinned with hand-computed RFC 5389 byte vectors — not round-tripped through our own
/// encoder, so an encode/decode bug can't cancel itself out. Plus one live loopback
/// round-trip: a miniature in-process STUN responder answers a real `StunClient.query`
/// over real UDP sockets.
@Suite struct StunTests {
    private let txid = Data([
        0xB7, 0xE7, 0xA7, 0x01, 0xBC, 0x34, 0xD6, 0x86, 0xFA, 0x87, 0xDF, 0xAE,
    ])

    @Test func bindingRequestWireShape() {
        let request = StunMessage.bindingRequest(transactionID: txid)
        #expect(request.count == 20)
        // Type 0x0001 (Binding Request), length 0, magic cookie, then our transaction ID.
        #expect(Data(request.prefix(8)) == Data([0x00, 0x01, 0x00, 0x00, 0x21, 0x12, 0xA4, 0x42]))
        #expect(Data(request.suffix(12)) == txid)
    }

    @Test func xorMappedAddressV4Decodes() {
        // 192.0.2.1:32853 → port 0x8055 ^ 0x2112 = 0xA147; addr 0xC0000201 ^ cookie = 0xE112A643.
        var response = Data([0x01, 0x01, 0x00, 0x0C, 0x21, 0x12, 0xA4, 0x42])
        response += txid
        response += Data([0x00, 0x20, 0x00, 0x08]) // XOR-MAPPED-ADDRESS, length 8
        response += Data([0x00, 0x01, 0xA1, 0x47, 0xE1, 0x12, 0xA6, 0x43])
        let mapped = StunMessage.parseMappedAddress(response, transactionID: txid)
        #expect(mapped == .v4(ip: "192.0.2.1", port: 32_853))
    }

    @Test func xorMappedAddressV6Decodes() {
        // 2001:db8::1 port 32853, XOR key = cookie ++ transaction ID.
        var response = Data([0x01, 0x01, 0x00, 0x18, 0x21, 0x12, 0xA4, 0x42])
        response += txid
        response += Data([0x00, 0x20, 0x00, 0x14]) // XOR-MAPPED-ADDRESS, length 20
        response += Data([0x00, 0x02, 0xA1, 0x47])
        response += Data([
            0x01, 0x13, 0xA9, 0xFA, 0xB7, 0xE7, 0xA7, 0x01,
            0xBC, 0x34, 0xD6, 0x86, 0xFA, 0x87, 0xDF, 0xAF,
        ])
        let mapped = StunMessage.parseMappedAddress(response, transactionID: txid)
        #expect(mapped == .v6(ip: "2001:db8::1", port: 32_853))
    }

    @Test func legacyMappedAddressDecodesWhenNoXORPresent() {
        // 203.0.113.9:4242, plain MAPPED-ADDRESS (0x0001), no XOR.
        var response = Data([0x01, 0x01, 0x00, 0x0C, 0x21, 0x12, 0xA4, 0x42])
        response += txid
        response += Data([0x00, 0x01, 0x00, 0x08])
        response += Data([0x00, 0x01, 0x10, 0x92, 0xCB, 0x00, 0x71, 0x09])
        let mapped = StunMessage.parseMappedAddress(response, transactionID: txid)
        #expect(mapped == .v4(ip: "203.0.113.9", port: 4_242))
    }

    @Test func rejectsWrongTransactionAndWrongType() {
        var response = Data([0x01, 0x01, 0x00, 0x0C, 0x21, 0x12, 0xA4, 0x42])
        response += Data(repeating: 0xFF, count: 12) // not our transaction
        response += Data([0x00, 0x20, 0x00, 0x08])
        response += Data([0x00, 0x01, 0xA1, 0x47, 0xE1, 0x12, 0xA6, 0x43])
        #expect(StunMessage.parseMappedAddress(response, transactionID: txid) == nil)

        let request = StunMessage.bindingRequest(transactionID: txid)
        #expect(StunMessage.parseMappedAddress(request, transactionID: txid) == nil)

        #expect(StunMessage.parseMappedAddress(Data([0x01, 0x01]), transactionID: txid) == nil)
    }

    @Test func skipsUnknownAttributesAndHonorsPadding() {
        // SOFTWARE (0x8022) with 5 bytes of value (3 padding) before the mapped address —
        // the scanner must land exactly on the next attribute boundary.
        var response = Data([0x01, 0x01, 0x00, 0x18, 0x21, 0x12, 0xA4, 0x42])
        response += txid
        response += Data([0x80, 0x22, 0x00, 0x05])
        response += Data([0x74, 0x65, 0x73, 0x74, 0x21, 0x00, 0x00, 0x00])
        response += Data([0x00, 0x20, 0x00, 0x08])
        response += Data([0x00, 0x01, 0xA1, 0x47, 0xE1, 0x12, 0xA6, 0x43])
        let mapped = StunMessage.parseMappedAddress(response, transactionID: txid)
        #expect(mapped == .v4(ip: "192.0.2.1", port: 32_853))
    }

    /// Full query over real sockets: a miniature in-process responder echoes the observed
    /// source as XOR-MAPPED-ADDRESS — exactly what a public STUN server does — and the
    /// client must decode its own loopback address.
    @Test func liveLoopbackQueryRoundTrips() async throws {
        let server = try UDPSocket(port: 0, ipv6: false)
        let client = try UDPSocket(port: 0, ipv6: false)
        defer {
            server.closeSocket()
            client.closeSocket()
        }
        let clock = ContinuousClock()

        let responder = Task {
            guard let packet = try await server.receive(deadline: clock.now + .seconds(5)) else {
                return
            }
            let requestTxid = Data(packet.data.suffix(12))
            var response = Data([0x01, 0x01, 0x00, 0x0C, 0x21, 0x12, 0xA4, 0x42])
            response += requestTxid
            response += Data([0x00, 0x20, 0x00, 0x08])
            let xport = packet.from.port ^ UInt16(StunMessage.magicCookie >> 16)
            response += Data([0x00, 0x01, UInt8(xport >> 8), UInt8(xport & 0xFF)])
            let octets = packet.from.ip.split(separator: ".").compactMap { UInt8($0) }
            let cookie: [UInt8] = [0x21, 0x12, 0xA4, 0x42]
            response += Data(zip(octets, cookie).map(^))
            try server.send(response, to: packet.from)
        }

        let mapped = await StunClient.query(
            socket: client,
            server: .v4(ip: "127.0.0.1", port: server.localPort),
            timeoutPerAttempt: .seconds(2), attempts: 2)
        try await responder.value
        #expect(mapped == .v4(ip: "127.0.0.1", port: client.localPort))
    }
}
