import Foundation

/// Big-endian byte writer for the wire formats (STUN, DTLS, RTP).
struct ByteWriter {
    private(set) var bytes: [UInt8] = []

    init(capacity: Int = 64) { bytes.reserveCapacity(capacity) }

    mutating func u8(_ v: UInt8) { bytes.append(v) }
    mutating func u16(_ v: UInt16) { bytes += [UInt8(v >> 8), UInt8(v & 0xFF)] }
    mutating func u24(_ v: Int) { bytes += [UInt8((v >> 16) & 0xFF), UInt8((v >> 8) & 0xFF), UInt8(v & 0xFF)] }
    mutating func u32(_ v: UInt32) { for shift in stride(from: 24, through: 0, by: -8) { bytes.append(UInt8((v >> UInt32(shift)) & 0xFF)) } }
    mutating func u48(_ v: UInt64) { for shift in stride(from: 40, through: 0, by: -8) { bytes.append(UInt8((v >> UInt64(shift)) & 0xFF)) } }
    mutating func u64(_ v: UInt64) { for shift in stride(from: 56, through: 0, by: -8) { bytes.append(UInt8((v >> UInt64(shift)) & 0xFF)) } }
    mutating func append<C: Collection>(_ data: C) where C.Element == UInt8 { bytes += data }

    /// Writes `body` prefixed by its length in `lengthBytes` bytes (1, 2 or 3).
    mutating func vector(_ lengthBytes: Int, _ body: (inout ByteWriter) -> Void) {
        var inner = ByteWriter()
        body(&inner)
        let n = inner.bytes.count
        switch lengthBytes {
        case 1: u8(UInt8(n))
        case 2: u16(UInt16(n))
        default: u24(n)
        }
        bytes += inner.bytes
    }
}

/// Big-endian reader that throws instead of trapping on short input.
struct ByteReader {
    struct OutOfBounds: Error {}

    let bytes: [UInt8]
    private(set) var offset: Int

    init<C: Collection>(_ bytes: C, offset: Int = 0) where C.Element == UInt8 {
        self.bytes = Array(bytes)
        self.offset = offset
    }

    var remaining: Int { bytes.count - offset }
    var isAtEnd: Bool { offset >= bytes.count }

    mutating func u8() throws -> UInt8 {
        guard remaining >= 1 else { throw OutOfBounds() }
        defer { offset += 1 }
        return bytes[offset]
    }

    mutating func u16() throws -> UInt16 { try UInt16(read(2)) }
    mutating func u24() throws -> Int { Int(try read(3)) }
    mutating func u32() throws -> UInt32 { try UInt32(read(4)) }
    mutating func u48() throws -> UInt64 { try read(6) }
    mutating func u64() throws -> UInt64 { try read(8) }

    mutating func take(_ n: Int) throws -> [UInt8] {
        guard n >= 0, remaining >= n else { throw OutOfBounds() }
        defer { offset += n }
        return Array(bytes[offset..<offset + n])
    }

    mutating func skip(_ n: Int) throws { _ = try take(n) }

    /// Reads a length-prefixed vector (`lengthBytes` = 1, 2 or 3).
    mutating func vector(_ lengthBytes: Int) throws -> [UInt8] {
        let n: Int = switch lengthBytes {
        case 1: Int(try u8())
        case 2: Int(try u16())
        default: try u24()
        }
        return try take(n)
    }

    private mutating func read(_ n: Int) throws -> UInt64 {
        guard remaining >= n else { throw OutOfBounds() }
        var v: UInt64 = 0
        for i in 0..<n { v = (v << 8) | UInt64(bytes[offset + i]) }
        offset += n
        return v
    }
}

extension Array where Element == UInt8 {
    init?(hex: String) {
        let clean = hex.filter { !$0.isWhitespace && $0 != ":" }
        guard clean.count.isMultiple(of: 2) else { return nil }
        var out: [UInt8] = []
        var index = clean.startIndex
        while index < clean.endIndex {
            let next = clean.index(index, offsetBy: 2)
            guard let b = UInt8(clean[index..<next], radix: 16) else { return nil }
            out.append(b)
            index = next
        }
        self = out
    }

    var hex: String { map { String(format: "%02x", $0) }.joined() }
}

/// Compares in constant time (MACs, Finished messages).
func constantTimeEqual(_ a: [UInt8], _ b: [UInt8]) -> Bool {
    guard a.count == b.count else { return false }
    var diff: UInt8 = 0
    for i in a.indices { diff |= a[i] ^ b[i] }
    return diff == 0
}