import CryptoKit
import Foundation
import Testing
@testable import Intercom

@Suite struct STUNTests {
    // RFC 5769 2.1: sample request, password "VOkJxbRl1RmTxUk/WvJxBt".
    let request = [UInt8](hex: """
        000100582112a442b7e7a701bc34d686fa87dfae
        802200105354554e207465737420636c69656e74
        002400046e0001ff
        80290008932ff9b151263b36
        000600096576746a3a68367659202020
        000800149aeaa70cbfd8cb56781ef2b5b2d3f249c1b571a2
        80280004e57a3bcf
        """)!
    let password = Array("VOkJxbRl1RmTxUk/WvJxBt".utf8)

    @Test func parsesAndVerifiesRFC5769Request() throws {
        let message = try #require(STUN.parse(request))
        #expect(message.type == STUN.bindingRequest)
        #expect(message.username == "evtj:h6vY")
        #expect(STUN.verify(request, message: message, key: password))
        #expect(!STUN.verify(request, message: message, key: Array("wrong".utf8)))
    }

    @Test func rejectsTamperedRequest() throws {
        var tampered = request
        tampered[45] ^= 1
        let message = try #require(STUN.parse(tampered))
        #expect(!STUN.verify(tampered, message: message, key: password))
    }

    @Test func buildsVerifiableResponse() throws {
        let tid = Array(request[8..<20])
        let response = STUN.bindingResponse(transactionID: tid, address: [192, 0, 2, 1], port: 32853, key: password)
        let message = try #require(STUN.parse(response))
        #expect(message.type == STUN.bindingSuccess)
        #expect(message.transactionID == tid)
        #expect(STUN.verify(response, message: message, key: password))
        // RFC 5769 2.2: XOR-MAPPED-ADDRESS for 192.0.2.1:32853.
        #expect(message.value(.xorMappedAddress) == [UInt8](hex: "0001a147e112a643"))
    }

    @Test func demultiplexesByFirstByte() {
        #expect(STUN.looksLikeSTUN(request))
        #expect(!STUN.looksLikeSTUN([22, 0xFE, 0xFD] + [UInt8](repeating: 0, count: 30)))
    }
}

@Suite struct KeyDerivationTests {
    @Test func tlsPRFMatchesOpenSSL() {
        let out = TLSPRF.compute(secret: [UInt8](hex: "9bbe436ba940f017b17652849a71db35")!, label: "test label",
                                 seed: [UInt8](hex: "a0ba9f936cda311827a6f796ffd5198c")!, length: 100)
        #expect(out.hex == """
            e3f229ba727be17b8d122620557cd453c2aab21d07c3d495329b52d4e61edb5a6b301791e90d35c9c9a46b4e14baf9af\
            0fa022f7077def17abfd3797c0564bab4fbc91666e9def9b97fce34f796789baa48082d122ee42c5a72e5a5110fff701\
            87347b66
            """)
    }

    // RFC 3711 B.2
    @Test func aesCounterKeystream() {
        let block = AESBlock(key: [UInt8](hex: "2B7E151628AED2A6ABF7158809CF4F3C")!)
        let stream = block.keystream(iv: [UInt8](hex: "F0F1F2F3F4F5F6F7F8F9FAFBFCFD0000")!, count: 48)
        #expect(stream.hex == "e03ead0935c95e80e166b16dd92b4eb4d23513162b02d0f72a43a2fe4a5f97ab41e95b3bb0a2e8dd477901e4fca894c0")
    }

    // RFC 3711 B.3
    @Test func srtpKeyDerivation() {
        let key = [UInt8](hex: "E1F97A0D3E018BE0D64FA32C06DE4139")!
        let salt = [UInt8](hex: "0EC675AD498AFEEBB6960B3AABE6")!
        #expect(SRTPContext.deriveKey(masterKey: key, masterSalt: salt, label: 0, length: 16).hex == "c61e7a93744f39ee10734afe3ff7a087")
        #expect(SRTPContext.deriveKey(masterKey: key, masterSalt: salt, label: 2, length: 14).hex == "30cbbc08863d8c85d49db34a9ae1")
        #expect(SRTPContext.deriveKey(masterKey: key, masterSalt: salt, label: 1, length: 20).hex == "cebe321f6ff7716b6fd4ab49af256a156d38baa4")
    }

    @Test func srtcpRoundTrips() throws {
        let sender = SRTPContext(masterKey: Array(repeating: 7, count: 16), masterSalt: Array(repeating: 9, count: 14))
        let receiver = SRTPContext(masterKey: Array(repeating: 7, count: 16), masterSalt: Array(repeating: 9, count: 14))
        // PLI: V=2, FMT=1, PT=206, length 2, sender SSRC, media SSRC
        let pli: [UInt8] = [0x81, 206, 0, 2, 1, 2, 3, 4, 5, 6, 7, 8]
        let protected = sender.protectRTCP(pli)
        #expect(protected != pli)
        #expect(receiver.unprotectRTCP(protected) == pli)
        var broken = protected
        broken[9] ^= 1
        #expect(receiver.unprotectRTCP(broken) == nil)
    }

    @Test func srtpEncryptsPayloadOnlyAndAppendsTag() {
        let context = SRTPContext(masterKey: Array(repeating: 1, count: 16), masterSalt: Array(repeating: 2, count: 14))
        let packet: [UInt8] = [0x80, 111, 0, 1, 0, 0, 0, 0, 0xAA, 0xBB, 0xCC, 0xDD] + Array(repeating: 0x55, count: 20)
        let out = context.protectRTP(packet)
        #expect(out.count == packet.count + 10)
        #expect(Array(out.prefix(12)) == Array(packet.prefix(12)))
        #expect(Array(out[12..<32]) != Array(packet[12...]))
    }
}

@Suite struct CertificateTests {
    @Test func selfSignedCertificateCarriesKeyAndValidSignature() throws {
        let cert = try DTLSCertificate()
        let spki = try #require(DER.subjectPublicKeyInfo(certificate: cert.der))
        let publicKey = try P256.Signing.PublicKey(derRepresentation: spki)
        #expect(publicKey.rawRepresentation == cert.privateKey.publicKey.rawRepresentation)

        // Signature over TBSCertificate verifies with the embedded key.
        let outer = try #require(DER.elements(cert.der, in: 0..<cert.der.count)?.first)
        let parts = try #require(DER.elements(cert.der, in: outer.contentRange))
        #expect(parts.count == 3)
        let tbs = Array(cert.der[parts[0].range])
        let signatureBits = Array(cert.der[parts[2].contentRange].dropFirst())
        let signature = try P256.Signing.ECDSASignature(derRepresentation: signatureBits)
        #expect(publicKey.isValidSignature(signature, for: tbs))
        #expect(cert.fingerprint.count == 32 * 3 - 1)
    }
}