import CryptoKit
import Foundation

/// TLS 1.2 PRF with SHA-256 (RFC 5246 section 5).
enum TLSPRF {
    static func compute(secret: [UInt8], label: String, seed: [UInt8], length: Int) -> [UInt8] {
        let key = SymmetricKey(data: secret)
        let fullSeed = Array(label.utf8) + seed
        var a = Array(HMAC<SHA256>.authenticationCode(for: fullSeed, using: key))
        var out: [UInt8] = []
        out.reserveCapacity(length + 32)
        while out.count < length {
            out += HMAC<SHA256>.authenticationCode(for: a + fullSeed, using: key)
            a = Array(HMAC<SHA256>.authenticationCode(for: a, using: key))
        }
        return Array(out.prefix(length))
    }
}

/// Minimal DER encoding/decoding for one self-signed ECDSA certificate and for reading peer keys.
enum DER {
    static func tlv(_ tag: UInt8, _ content: [UInt8]) -> [UInt8] {
        var out = [tag]
        let n = content.count
        if n < 0x80 {
            out.append(UInt8(n))
        } else if n <= 0xFF {
            out += [0x81, UInt8(n)]
        } else {
            out += [0x82, UInt8(n >> 8), UInt8(n & 0xFF)]
        }
        return out + content
    }

    static func sequence(_ items: [UInt8]...) -> [UInt8] { tlv(0x30, items.flatMap { $0 }) }
    static func set(_ items: [UInt8]...) -> [UInt8] { tlv(0x31, items.flatMap { $0 }) }

    static func integer(_ bytes: [UInt8]) -> [UInt8] {
        var b = Array(bytes.drop { $0 == 0 })
        if b.isEmpty { b = [0] }
        if b[0] & 0x80 != 0 { b.insert(0, at: 0) }
        return tlv(0x02, b)
    }

    static func oid(_ components: [UInt]) -> [UInt8] {
        var body = [UInt8(components[0] * 40 + components[1])]
        for c in components.dropFirst(2) {
            var chunk = [UInt8(c & 0x7F)]
            var v = c >> 7
            while v > 0 {
                chunk.insert(UInt8(v & 0x7F) | 0x80, at: 0)
                v >>= 7
            }
            body += chunk
        }
        return tlv(0x06, body)
    }

    static func utf8(_ s: String) -> [UInt8] { tlv(0x0C, Array(s.utf8)) }
    static func bitString(_ bytes: [UInt8]) -> [UInt8] { tlv(0x03, [0] + bytes) }

    static func utcTime(_ date: Date) -> [UInt8] {
        let f = DateFormatter()
        f.locale = Locale(identifier: "en_US_POSIX")
        f.timeZone = TimeZone(identifier: "UTC")
        f.dateFormat = "yyMMddHHmmss'Z'"
        return tlv(0x17, Array(f.string(from: date).utf8))
    }

    /// One parsed element: tag and the full encoding range inside the source buffer.
    struct Element {
        let tag: UInt8
        let header: Int
        let range: Range<Int>
        var contentRange: Range<Int> { range.lowerBound + header..<range.upperBound }
    }

    static func elements(_ bytes: [UInt8], in range: Range<Int>) -> [Element]? {
        var out: [Element] = []
        var i = range.lowerBound
        while i < range.upperBound {
            guard i + 2 <= range.upperBound else { return nil }
            let tag = bytes[i]
            var length = Int(bytes[i + 1])
            var header = 2
            if length & 0x80 != 0 {
                let count = length & 0x7F
                guard count <= 3, i + 2 + count <= range.upperBound else { return nil }
                length = 0
                for k in 0..<count { length = length << 8 | Int(bytes[i + 2 + k]) }
                header += count
            }
            guard i + header + length <= range.upperBound else { return nil }
            out.append(Element(tag: tag, header: header, range: i..<i + header + length))
            i += header + length
        }
        return out
    }

    /// SubjectPublicKeyInfo of an X.509 certificate in DER.
    static func subjectPublicKeyInfo(certificate: [UInt8]) -> [UInt8]? {
        guard let cert = elements(certificate, in: 0..<certificate.count)?.first, cert.tag == 0x30,
              let tbs = elements(certificate, in: cert.contentRange)?.first, tbs.tag == 0x30,
              let fields = elements(certificate, in: tbs.contentRange) else { return nil }
        // [0] version (optional), serial, signature, issuer, validity, subject, subjectPublicKeyInfo
        let start = fields.first?.tag == 0xA0 ? 1 : 0
        guard fields.count > start + 5 else { return nil }
        return Array(certificate[fields[start + 5].range])
    }
}

/// The self-signed ECDSA P-256 certificate that identifies this Musubi in DTLS.
/// Browsers only check it against the fingerprint in the SDP, not against a CA.
struct DTLSCertificate: Sendable {
    let privateKey: P256.Signing.PrivateKey
    let der: [UInt8]

    /// SHA-256 fingerprint as written in SDP (`AB:CD:…`).
    var fingerprint: String { Self.fingerprint(der) }

    static func fingerprint(_ der: [UInt8]) -> String {
        SHA256.hash(data: der).map { String(format: "%02X", $0) }.joined(separator: ":")
    }

    init(commonName: String = "Musubi", now: Date = Date()) throws {
        let key = P256.Signing.PrivateKey()
        let ecdsaWithSHA256 = DER.sequence(DER.oid([1, 2, 840, 10045, 4, 3, 2]))
        let name = DER.sequence(DER.set(DER.sequence(DER.oid([2, 5, 4, 3]), DER.utf8(commonName))))
        let validity = DER.sequence(DER.utcTime(now.addingTimeInterval(-86_400)), DER.utcTime(now.addingTimeInterval(30 * 86_400)))
        let spki = DER.sequence(
            DER.sequence(DER.oid([1, 2, 840, 10045, 2, 1]), DER.oid([1, 2, 840, 10045, 3, 1, 7])),
            DER.bitString(Array(key.publicKey.x963Representation)))
        var serial = [UInt8](repeating: 0, count: 8)
        _ = SecRandomCopyBytes(kSecRandomDefault, serial.count, &serial)
        serial[0] &= 0x7F
        let tbs = DER.sequence(
            DER.tlv(0xA0, DER.integer([2])),
            DER.integer(serial),
            ecdsaWithSHA256, name, validity, name, spki)
        let signature = try key.signature(for: tbs).derRepresentation
        self.privateKey = key
        self.der = DER.sequence(tbs, ecdsaWithSHA256, DER.bitString(Array(signature)))
    }
}