import CryptoKit
import Foundation

/// The subset of STUN (RFC 5389) an ICE-lite agent needs: answer Binding requests
/// with short-term credentials (RFC 8445), MESSAGE-INTEGRITY and FINGERPRINT.
enum STUN {
    static let magicCookie: UInt32 = 0x2112_A442
    static let bindingRequest: UInt16 = 0x0001
    static let bindingSuccess: UInt16 = 0x0101

    enum Attribute: UInt16 {
        case username = 0x0006
        case messageIntegrity = 0x0008
        case xorMappedAddress = 0x0020
        case priority = 0x0024
        case useCandidate = 0x0025
        case fingerprint = 0x8028
        case iceControlled = 0x8029
        case iceControlling = 0x802A
    }

    struct Message {
        var type: UInt16
        var transactionID: [UInt8]
        /// Attributes in wire order, with the offset of each attribute header in the packet.
        var attributes: [(type: UInt16, value: [UInt8], offset: Int)]

        func value(_ attribute: Attribute) -> [UInt8]? {
            attributes.first { $0.type == attribute.rawValue }?.value
        }

        var username: String? { value(.username).flatMap { String(bytes: $0, encoding: .utf8) } }
    }

    /// RFC 7983 demultiplexing: first byte 0–3 and the magic cookie.
    static func looksLikeSTUN(_ packet: [UInt8]) -> Bool {
        packet.count >= 20 && packet[0] < 4
            && packet[4] == 0x21 && packet[5] == 0x12 && packet[6] == 0xA4 && packet[7] == 0x42
    }

    static func parse(_ packet: [UInt8]) -> Message? {
        guard looksLikeSTUN(packet) else { return nil }
        var reader = ByteReader(packet)
        do {
            let type = try reader.u16()
            let length = Int(try reader.u16())
            guard length + 20 == packet.count, length.isMultiple(of: 4) else { return nil }
            try reader.skip(4)
            let transactionID = try reader.take(12)
            var attributes: [(UInt16, [UInt8], Int)] = []
            while !reader.isAtEnd {
                let offset = reader.offset
                let attrType = try reader.u16()
                let attrLength = Int(try reader.u16())
                let value = try reader.take(attrLength)
                try reader.skip((4 - attrLength % 4) % 4)
                attributes.append((attrType, value, offset))
            }
            return Message(type: type, transactionID: transactionID, attributes: attributes)
        } catch {
            return nil
        }
    }

    /// Checks MESSAGE-INTEGRITY (HMAC-SHA1 with `key`) and, if present, FINGERPRINT.
    static func verify(_ packet: [UInt8], message: Message, key: [UInt8]) -> Bool {
        if let fp = message.attributes.first(where: { $0.type == Attribute.fingerprint.rawValue }) {
            guard fp.value == fingerprint(packet, upTo: fp.offset) else { return false }
        }
        guard let mi = message.attributes.first(where: { $0.type == Attribute.messageIntegrity.rawValue }) else { return false }
        return constantTimeEqual(mi.value, integrity(packet, upTo: mi.offset, key: key))
    }

    /// Builds a Binding success response with XOR-MAPPED-ADDRESS, MESSAGE-INTEGRITY and FINGERPRINT.
    static func bindingResponse(transactionID: [UInt8], address: [UInt8], port: UInt16, key: [UInt8]) -> [UInt8] {
        var w = ByteWriter(capacity: 80)
        w.u16(bindingSuccess)
        w.u16(0) // length, patched below
        w.u32(magicCookie)
        w.append(transactionID)

        let cookie: [UInt8] = [0x21, 0x12, 0xA4, 0x42]
        w.u16(Attribute.xorMappedAddress.rawValue)
        w.u16(address.count == 4 ? 8 : 20)
        w.u8(0)
        w.u8(address.count == 4 ? 0x01 : 0x02)
        w.u16(port ^ 0x2112)
        let mask = cookie + transactionID
        w.append(address.enumerated().map { $0.element ^ mask[$0.offset] })

        var bytes = w.bytes
        let miOffset = bytes.count
        let mac = integrity(bytes, upTo: miOffset, key: key)
        var tail = ByteWriter()
        tail.u16(Attribute.messageIntegrity.rawValue)
        tail.u16(20)
        tail.append(mac)
        bytes += tail.bytes

        let fpOffset = bytes.count
        let crc = fingerprint(bytes, upTo: fpOffset)
        bytes += [0x80, 0x28, 0x00, 0x04] + crc
        setLength(&bytes, bytes.count - 20)
        return bytes
    }

    /// HMAC-SHA1 over the message up to `offset`, with the header length covering the MESSAGE-INTEGRITY attribute.
    private static func integrity(_ packet: [UInt8], upTo offset: Int, key: [UInt8]) -> [UInt8] {
        var prefix = Array(packet[0..<offset])
        setLength(&prefix, offset - 20 + 24)
        let mac = HMAC<Insecure.SHA1>.authenticationCode(for: prefix, using: SymmetricKey(data: key))
        return Array(mac)
    }

    /// CRC-32 over the message up to `offset`, length covering FINGERPRINT, XOR 0x5354554E.
    private static func fingerprint(_ packet: [UInt8], upTo offset: Int) -> [UInt8] {
        var prefix = Array(packet[0..<offset])
        setLength(&prefix, offset - 20 + 8)
        let value = CRC32.checksum(prefix) ^ 0x5354_554E
        return [UInt8(value >> 24), UInt8((value >> 16) & 0xFF), UInt8((value >> 8) & 0xFF), UInt8(value & 0xFF)]
    }

    private static func setLength(_ bytes: inout [UInt8], _ length: Int) {
        bytes[2] = UInt8(length >> 8)
        bytes[3] = UInt8(length & 0xFF)
    }
}

enum CRC32 {
    private static let table: [UInt32] = (0..<256).map { n in
        var c = UInt32(n)
        for _ in 0..<8 { c = (c & 1) != 0 ? 0xEDB8_8320 ^ (c >> 1) : c >> 1 }
        return c
    }

    static func checksum(_ bytes: [UInt8]) -> UInt32 {
        var crc: UInt32 = 0xFFFF_FFFF
        for b in bytes { crc = table[Int((crc ^ UInt32(b)) & 0xFF)] ^ (crc >> 8) }
        return crc ^ 0xFFFF_FFFF
    }
}