import CryptoKit
import Foundation

/// Keying material for SRTP, exported from DTLS (RFC 5764 4.2).
struct SRTPKeyMaterial: Sendable, Equatable {
    let clientKey: [UInt8]
    let serverKey: [UInt8]
    let clientSalt: [UInt8]
    let serverSalt: [UInt8]
}

/// A DTLS 1.2 server for WebRTC, deliberately narrow:
/// one cipher suite (TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256), P-256, ECDSA client certificates
/// checked against the SDP fingerprint, the use_srtp extension with SRTP_AES128_CM_HMAC_SHA1_80
/// and extended master secret. No resumption, no renegotiation.
/// Synchronous state machine: feed datagrams, send what comes back. Not thread-safe.
final class DTLSServer {
    enum State: Equatable {
        case awaitingClientHello
        case awaitingClientFlight
        case established
        case closed
        case failed(String)
    }

    struct Failure: Error, CustomStringConvertible {
        let description: String
        let alert: UInt8
        init(_ description: String, alert: UInt8 = 40) { self.description = description; self.alert = alert }
    }

    private enum ContentType: UInt8 { case changeCipherSpec = 20, alert = 21, handshake = 22, applicationData = 23 }
    private enum HandshakeType: UInt8 {
        case clientHello = 1, serverHello = 2, certificate = 11, serverKeyExchange = 12, certificateRequest = 13
        case serverHelloDone = 14, certificateVerify = 15, clientKeyExchange = 16, finished = 20
    }

    static let cipherSuite: UInt16 = 0xC02B
    static let srtpProfile: UInt16 = 0x0001
    private static let version: UInt16 = 0xFEFD
    private static let ecdsaSecp256r1SHA256: UInt16 = 0x0403

    private(set) var state: State = .awaitingClientHello
    private(set) var keyMaterial: SRTPKeyMaterial?

    private let certificate: DTLSCertificate
    private let expectedFingerprint: String?

    private var transcript: [UInt8] = []
    private var clientRandom: [UInt8] = []
    private var serverRandom: [UInt8] = []
    private var usesExtendedMasterSecret = false
    private var echoRenegotiationInfo = false
    private var echoPointFormats = false
    private var ephemeral = P256.KeyAgreement.PrivateKey()
    private var masterSecret: [UInt8] = []
    private var clientPublicKey: P256.Signing.PublicKey?

    private var readCipher: (key: SymmetricKey, iv: [UInt8])?
    private var writeCipher: (key: SymmetricKey, iv: [UInt8])?
    private var writeEpoch: UInt16 = 0
    private var writeSequence: [UInt16: UInt64] = [:]
    private var nextSendMessageSeq: UInt16 = 0
    private var nextReceiveMessageSeq: UInt16 = 0
    private var fragments: [UInt16: (type: UInt8, body: [UInt8?])] = [:]
    private var lastFlight: [(type: ContentType, epoch: UInt16, body: [UInt8])] = []
    private var flightTrigger: UInt16 = 0

    /// `expectedFingerprint`: the peer's `a=fingerprint:sha-256` from the SDP offer.
    init(certificate: DTLSCertificate, expectedFingerprint: String?) {
        self.certificate = certificate
        self.expectedFingerprint = expectedFingerprint?.uppercased()
    }

    /// RFC 7983 demultiplexing: DTLS content types 20…63.
    static func looksLikeDTLS(_ packet: [UInt8]) -> Bool {
        guard let first = packet.first else { return false }
        return first >= 20 && first <= 63
    }

    /// Processes one datagram and returns the datagrams to send back.
    func handle(_ datagram: [UInt8]) -> [[UInt8]] {
        switch state {
        case .closed, .failed: return []
        default: break
        }
        do {
            var output: [[UInt8]] = []
            var reader = ByteReader(datagram)
            while reader.remaining >= 13 {
                let type = try reader.u8()
                _ = try reader.u16()
                let epoch = try reader.u16()
                let sequence = try reader.u48()
                let fragment = try reader.take(Int(try reader.u16()))
                guard let content = ContentType(rawValue: type) else { continue }
                let plain: [UInt8]
                if epoch == 0 {
                    plain = fragment
                } else if epoch == 1, readCipher != nil {
                    guard let decrypted = decrypt(fragment, type: type, epoch: epoch, sequence: sequence) else { continue }
                    plain = decrypted
                } else {
                    continue
                }
                output += try process(content, epoch: epoch, plain)
            }
            return output
        } catch let failure as Failure {
            state = .failed(failure.description)
            return [record(.alert, epoch: writeEpoch, body: [2, failure.alert])]
        } catch {
            state = .failed("Malformed DTLS record")
            return [record(.alert, epoch: writeEpoch, body: [2, 50])]
        }
    }

    // MARK: - Records

    private func process(_ type: ContentType, epoch: UInt16, _ body: [UInt8]) throws -> [[UInt8]] {
        switch type {
        case .handshake:
            return try receiveHandshake(body, epoch: epoch)
        case .changeCipherSpec:
            if state == .awaitingClientFlight, readCipher == nil, !masterSecret.isEmpty {
                readCipher = keys().client
            }
            return []
        case .alert:
            if body.count >= 2, body[1] == 0 { state = .closed } else if body.first == 2 { state = .failed("Peer alert \(body.last ?? 0)") }
            return []
        case .applicationData:
            return []
        }
    }

    private func receiveHandshake(_ body: [UInt8], epoch: UInt16) throws -> [[UInt8]] {
        var reader = ByteReader(body)
        var output: [[UInt8]] = []
        while reader.remaining >= 12 {
            let type = try reader.u8()
            let length = try reader.u24()
            let messageSeq = try reader.u16()
            let offset = try reader.u24()
            let fragmentLength = try reader.u24()
            let fragment = try reader.take(fragmentLength)
            guard length <= 1 << 16, offset + fragmentLength <= length else { throw Failure("Bad fragment", alert: 50) }

            if messageSeq < nextReceiveMessageSeq {
                // Retransmitted flight: the peer missed ours.
                if messageSeq == flightTrigger, !lastFlight.isEmpty { output += resendLastFlight() }
                continue
            }
            var entry = fragments[messageSeq] ?? (type, [UInt8?](repeating: nil, count: length))
            guard entry.type == type, entry.body.count == length else { throw Failure("Inconsistent fragments", alert: 50) }
            for i in 0..<fragmentLength { entry.body[offset + i] = fragment[i] }
            fragments[messageSeq] = entry

            while let next = fragments[nextReceiveMessageSeq], !next.body.contains(where: { $0 == nil }) {
                fragments[nextReceiveMessageSeq] = nil
                let seq = nextReceiveMessageSeq
                nextReceiveMessageSeq += 1
                output += try handleMessage(next.type, next.body.map { $0! }, messageSeq: seq, epoch: epoch)
            }
        }
        return output
    }

    private func handleMessage(_ type: UInt8, _ body: [UInt8], messageSeq: UInt16, epoch: UInt16) throws -> [[UInt8]] {
        let full = handshakeHeader(type: type, length: body.count, messageSeq: messageSeq) + body
        switch (state, HandshakeType(rawValue: type)) {
        case (.awaitingClientHello, .clientHello):
            try parseClientHello(body)
            transcript += full
            flightTrigger = messageSeq
            return sendServerFlight()

        case (.awaitingClientFlight, .certificate):
            var r = ByteReader(body)
            var list = ByteReader(try r.vector(3))
            guard !list.isAtEnd else { throw Failure("Client sent no certificate", alert: 42) }
            let der = try list.vector(3)
            if let expectedFingerprint, DTLSCertificate.fingerprint(der) != expectedFingerprint {
                throw Failure("Client certificate does not match the SDP fingerprint", alert: 42)
            }
            guard let spki = DER.subjectPublicKeyInfo(certificate: der),
                  let key = try? P256.Signing.PublicKey(derRepresentation: spki) else {
                throw Failure("Unsupported client certificate", alert: 43)
            }
            clientPublicKey = key
            transcript += full

        case (.awaitingClientFlight, .clientKeyExchange):
            var r = ByteReader(body)
            let point = try r.vector(1)
            guard let peer = try? P256.KeyAgreement.PublicKey(x963Representation: point),
                  let shared = try? ephemeral.sharedSecretFromKeyAgreement(with: peer) else {
                throw Failure("Bad client key share", alert: 47)
            }
            let preMaster = shared.withUnsafeBytes { Array($0) }
            transcript += full
            if usesExtendedMasterSecret {
                masterSecret = TLSPRF.compute(secret: preMaster, label: "extended master secret",
                                              seed: Array(SHA256.hash(data: transcript)), length: 48)
            } else {
                masterSecret = TLSPRF.compute(secret: preMaster, label: "master secret",
                                              seed: clientRandom + serverRandom, length: 48)
            }

        case (.awaitingClientFlight, .certificateVerify):
            guard let clientPublicKey, !masterSecret.isEmpty else { throw Failure("Unexpected CertificateVerify", alert: 10) }
            var r = ByteReader(body)
            guard try r.u16() == Self.ecdsaSecp256r1SHA256 else { throw Failure("Unsupported signature algorithm", alert: 40) }
            let signature = try r.vector(2)
            guard let parsed = try? P256.Signing.ECDSASignature(derRepresentation: signature),
                  clientPublicKey.isValidSignature(parsed, for: transcript) else {
                throw Failure("CertificateVerify failed", alert: 51)
            }
            transcript += full

        case (.awaitingClientFlight, .finished):
            guard epoch == 1, clientPublicKey != nil else { throw Failure("Finished before authentication", alert: 10) }
            let expected = TLSPRF.compute(secret: masterSecret, label: "client finished",
                                          seed: Array(SHA256.hash(data: transcript)), length: 12)
            guard constantTimeEqual(body, expected) else { throw Failure("Client Finished does not verify", alert: 51) }
            transcript += full
            flightTrigger = messageSeq
            let output = sendFinalFlight()
            exportKeys()
            state = .established
            return output

        default:
            throw Failure("Unexpected handshake message \(type)", alert: 10)
        }
        return []
    }

    // MARK: - ClientHello

    private func parseClientHello(_ body: [UInt8]) throws {
        var r = ByteReader(body)
        guard try r.u16() == Self.version else { throw Failure("Only DTLS 1.2 is supported", alert: 70) }
        clientRandom = try r.take(32)
        _ = try r.vector(1) // session id
        _ = try r.vector(1) // cookie
        var suites = ByteReader(try r.vector(2))
        var offersSuite = false
        while suites.remaining >= 2 {
            let suite = try suites.u16()
            if suite == Self.cipherSuite { offersSuite = true }
            if suite == 0x00FF { echoRenegotiationInfo = true }
        }
        guard offersSuite else { throw Failure("Client does not offer ECDHE-ECDSA-AES128-GCM-SHA256") }
        _ = try r.vector(1) // compression methods

        var offersSRTP = false
        var offersP256 = false
        if !r.isAtEnd {
            var extensions = ByteReader(try r.vector(2))
            while extensions.remaining >= 4 {
                let type = try extensions.u16()
                var data = ByteReader(try extensions.vector(2))
                switch type {
                case 0x000E: // use_srtp
                    var profiles = ByteReader(try data.vector(2))
                    while profiles.remaining >= 2 { if try profiles.u16() == Self.srtpProfile { offersSRTP = true } }
                case 0x000A: // supported_groups
                    var groups = ByteReader(try data.vector(2))
                    while groups.remaining >= 2 { if try groups.u16() == 23 { offersP256 = true } }
                case 0x0017: usesExtendedMasterSecret = true
                case 0xFF01: echoRenegotiationInfo = true
                case 0x000B: echoPointFormats = true
                default: break
                }
            }
        } else {
            offersP256 = true // no supported_groups extension means "any"
        }
        guard offersSRTP else { throw Failure("Client does not offer SRTP_AES128_CM_HMAC_SHA1_80") }
        guard offersP256 else { throw Failure("Client does not offer P-256") }
    }

    // MARK: - Server flights

    private func sendServerFlight() -> [[UInt8]] {
        serverRandom = (0..<32).map { _ in UInt8.random(in: 0...255) }

        var hello = ByteWriter()
        hello.u16(Self.version)
        hello.append(serverRandom)
        hello.u8(0) // session id
        hello.u16(Self.cipherSuite)
        hello.u8(0) // no compression
        hello.vector(2) { ext in
            if echoRenegotiationInfo { ext.u16(0xFF01); ext.u16(1); ext.u8(0) }
            if usesExtendedMasterSecret { ext.u16(0x0017); ext.u16(0) }
            if echoPointFormats { ext.u16(0x000B); ext.u16(2); ext.u8(1); ext.u8(0) }
            ext.u16(0x000E); ext.u16(5); ext.u16(2); ext.u16(Self.srtpProfile); ext.u8(0)
        }

        var cert = ByteWriter()
        cert.vector(3) { list in list.vector(3) { $0.append(certificate.der) } }

        var params = ByteWriter()
        params.u8(3) // named_curve
        params.u16(23) // secp256r1
        params.vector(1) { $0.append(ephemeral.publicKey.x963Representation) }
        let signed = (try? certificate.privateKey.signature(for: clientRandom + serverRandom + params.bytes).derRepresentation) ?? Data()
        var keyExchange = params
        keyExchange.u16(Self.ecdsaSecp256r1SHA256)
        keyExchange.vector(2) { $0.append(signed) }

        var request = ByteWriter()
        request.vector(1) { $0.u8(64) } // ecdsa_sign
        request.vector(2) { $0.u16(Self.ecdsaSecp256r1SHA256) }
        request.u16(0) // no CA names

        lastFlight = []
        for (type, body) in [(HandshakeType.serverHello, hello.bytes), (.certificate, cert.bytes),
                             (.serverKeyExchange, keyExchange.bytes), (.certificateRequest, request.bytes),
                             (.serverHelloDone, [])] {
            let message = handshakeHeader(type: type.rawValue, length: body.count, messageSeq: nextSendMessageSeq) + body
            nextSendMessageSeq += 1
            transcript += message
            lastFlight.append((.handshake, 0, message))
        }
        state = .awaitingClientFlight
        return resendLastFlight()
    }

    private func sendFinalFlight() -> [[UInt8]] {
        let verify = TLSPRF.compute(secret: masterSecret, label: "server finished",
                                    seed: Array(SHA256.hash(data: transcript)), length: 12)
        let finished = handshakeHeader(type: HandshakeType.finished.rawValue, length: 12, messageSeq: nextSendMessageSeq) + verify
        nextSendMessageSeq += 1
        transcript += finished
        writeCipher = keys().server
        lastFlight = [(.changeCipherSpec, 0, [1]), (.handshake, 1, finished)]
        return resendLastFlight()
    }

    /// Sends the last flight again with fresh record sequence numbers, all in one datagram.
    private func resendLastFlight() -> [[UInt8]] {
        [lastFlight.flatMap { record($0.type, epoch: $0.epoch, body: $0.body) }]
    }

    private func handshakeHeader(type: UInt8, length: Int, messageSeq: UInt16) -> [UInt8] {
        var w = ByteWriter(capacity: 12)
        w.u8(type)
        w.u24(length)
        w.u16(messageSeq)
        w.u24(0)
        w.u24(length)
        return w.bytes
    }

    private func record(_ type: ContentType, epoch: UInt16, body: [UInt8]) -> [UInt8] {
        let sequence = writeSequence[epoch, default: 0]
        writeSequence[epoch] = sequence + 1
        let payload = epoch == 0 ? body : encrypt(body, type: type.rawValue, epoch: epoch, sequence: sequence)
        var w = ByteWriter(capacity: 13 + payload.count)
        w.u8(type.rawValue)
        w.u16(Self.version)
        w.u16(epoch)
        w.u48(sequence)
        w.u16(UInt16(payload.count))
        w.append(payload)
        return w.bytes
    }

    // MARK: - Keys and AEAD

    private func keys() -> (client: (key: SymmetricKey, iv: [UInt8]), server: (key: SymmetricKey, iv: [UInt8])) {
        let block = TLSPRF.compute(secret: masterSecret, label: "key expansion", seed: serverRandom + clientRandom, length: 40)
        return ((SymmetricKey(data: block[0..<16]), Array(block[32..<36])),
                (SymmetricKey(data: block[16..<32]), Array(block[36..<40])))
    }

    private func exportKeys() {
        let m = TLSPRF.compute(secret: masterSecret, label: "EXTRACTOR-dtls_srtp", seed: clientRandom + serverRandom, length: 60)
        keyMaterial = SRTPKeyMaterial(clientKey: Array(m[0..<16]), serverKey: Array(m[16..<32]),
                                      clientSalt: Array(m[32..<46]), serverSalt: Array(m[46..<60]))
    }

    private func additionalData(type: UInt8, epoch: UInt16, sequence: UInt64, length: Int) -> [UInt8] {
        var w = ByteWriter(capacity: 13)
        w.u16(epoch)
        w.u48(sequence)
        w.u8(type)
        w.u16(Self.version)
        w.u16(UInt16(length))
        return w.bytes
    }

    private func encrypt(_ plain: [UInt8], type: UInt8, epoch: UInt16, sequence: UInt64) -> [UInt8] {
        guard let writeCipher else { return plain }
        var explicit = ByteWriter(capacity: 8)
        explicit.u16(epoch)
        explicit.u48(sequence)
        guard let nonce = try? AES.GCM.Nonce(data: writeCipher.iv + explicit.bytes),
              let sealed = try? AES.GCM.seal(plain, using: writeCipher.key, nonce: nonce,
                                             authenticating: additionalData(type: type, epoch: epoch, sequence: sequence, length: plain.count))
        else { return [] }
        return explicit.bytes + Array(sealed.ciphertext) + Array(sealed.tag)
    }

    private func decrypt(_ fragment: [UInt8], type: UInt8, epoch: UInt16, sequence: UInt64) -> [UInt8]? {
        guard let readCipher, fragment.count >= 8 + 16 else { return nil }
        let explicit = Array(fragment[0..<8])
        let ciphertext = Array(fragment[8..<fragment.count - 16])
        let tag = Array(fragment[(fragment.count - 16)...])
        guard let nonce = try? AES.GCM.Nonce(data: readCipher.iv + explicit),
              let box = try? AES.GCM.SealedBox(nonce: nonce, ciphertext: ciphertext, tag: tag),
              let plain = try? AES.GCM.open(box, using: readCipher.key,
                                            authenticating: additionalData(type: type, epoch: epoch, sequence: sequence, length: ciphertext.count))
        else { return nil }
        return Array(plain)
    }
}