Musubi/Packages/MusubiKit/Sources/Intercom/IntercomEngine.swift

import Foundation
import Network
import OSLog

/// Media transport for all phones: one UDP port on the clinic interface, demultiplexed into sessions.
/// Everything runs on one serial queue; public methods may be called from any thread.
final class IntercomEngine: @unchecked Sendable {
    struct SessionInfo: Sendable, Equatable {
        let id: UUID
        let name: String
        let role: String
        let connected: Bool
    }

    var onSessionsChanged: (@Sendable ([SessionInfo]) -> Void)?
    var onKeyFrameNeeded: (@Sendable () -> Void)?

    let queue = DispatchQueue(label: "studio.zielke.musubi.intercom.media", qos: .userInteractive)
    private let certificate: DTLSCertificate
    private let log = Logger(subsystem: "studio.zielke.musubi", category: "intercom")

    // queue only
    private var listener: NWListener?
    private var sessions: [UUID: PeerSession] = [:]
    private var connections: [UUID: NWConnection] = [:]
    private var timer: DispatchSourceTimer?

    init(certificate: DTLSCertificate) {
        self.certificate = certificate
    }

    /// Binds the media port; returns once listening (or throws).
    func start(interface: NWInterface, port: UInt16) async throws {
        let parameters = NWParameters.udp
        parameters.requiredInterface = interface
        parameters.allowLocalEndpointReuse = true
        let listener = try NWListener(using: parameters, on: NWEndpoint.Port(rawValue: port) ?? .any)
        let once = OnceFlag()
        try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
            listener.stateUpdateHandler = { [log] state in
                switch state {
                case .ready:
                    if once.claim() { continuation.resume() }
                case .failed(let error):
                    log.error("Media port failed: \(error.localizedDescription, privacy: .public)")
                    if once.claim() { continuation.resume(throwing: error) }
                default: break
                }
            }
            listener.newConnectionHandler = { [weak self] connection in self?.accept(connection) }
            listener.start(queue: queue)
            queue.async { self.listener = listener }
        }
        queue.async { self.startTimer() }
    }

    func stop() {
        queue.sync {
            timer?.cancel()
            timer = nil
            listener?.cancel()
            listener = nil
            connections.values.forEach { $0.cancel() }
            connections = [:]
            sessions = [:]
        }
        onSessionsChanged?([])
    }

    /// Creates a session for a browser offer and returns its id and the SDP answer.
    func createSession(offer: String, name: String, role: String, address: String, port: UInt16) throws -> (UUID, String)? {
        let accepted = try PeerSession.accept(offerSDP: offer, name: name, role: role, certificate: certificate,
                                              address: address, port: port)
        guard let accepted else { return nil }
        queue.sync {
            sessions[accepted.session.id] = accepted.session
            publish()
        }
        log.info("Session for \(name, privacy: .public) (\(role, privacy: .public)) created")
        return (accepted.session.id, accepted.answer)
    }

    func removeSession(_ id: UUID) {
        queue.async {
            self.sessions[id] = nil
            self.connections.removeValue(forKey: id)?.cancel()
            self.publish()
        }
    }

    /// Sends one Opus packet (20 ms) to every connected phone.
    func sendAudio(opus: [UInt8], timestamp: UInt32) {
        queue.async {
            for (id, session) in self.sessions {
                guard let packet = session.audioPacket(opus: opus, timestamp: timestamp) else { continue }
                self.connections[id]?.send(content: Data(packet), completion: .idempotent)
            }
        }
    }

    /// Sends one H.264 access unit to every connected phone.
    func sendVideo(nalUnits: [[UInt8]], timestamp: UInt32) {
        let payloads = H264Packetizer.payloads(for: nalUnits)
        queue.async {
            for (id, session) in self.sessions {
                guard let connection = self.connections[id] else { continue }
                for packet in session.videoPackets(payloads: payloads, timestamp: timestamp) {
                    connection.send(content: Data(packet), completion: .idempotent)
                }
            }
        }
    }

    // MARK: - queue

    private func accept(_ connection: NWConnection) {
        connection.start(queue: queue)
        receive(on: connection)
    }

    private func receive(on connection: NWConnection) {
        connection.receiveMessage { [weak self] data, _, _, error in
            guard let self else { return }
            if let data, !data.isEmpty { self.handle(Array(data), from: connection) }
            if error == nil { self.receive(on: connection) } else { connection.cancel() }
        }
    }

    private func handle(_ packet: [UInt8], from connection: NWConnection) {
        switch PeerSession.classify(packet) {
        case .stun:
            guard let message = STUN.parse(packet),
                  case .hostPort(let host, let port) = connection.endpoint else { return }
            let address: [UInt8]
            switch host {
            case .ipv4(let v4): address = Array(v4.rawValue)
            case .ipv6(let v6):
                let raw = Array(v6.rawValue)
                // IPv4-mapped (::ffff:a.b.c.d) on a dual-stack socket.
                address = raw.prefix(12).elementsEqual([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xFF, 0xFF]) ? Array(raw.suffix(4)) : raw
            default: return
            }
            let localUfrag = message.username?.split(separator: ":").first.map(String.init)
            guard let session = sessions.values.first(where: { $0.localUfrag == localUfrag }),
                  let reply = session.answerSTUN(packet, message: message, address: address, port: port.rawValue) else { return }
            if connections[session.id] !== connection {
                connections[session.id] = connection
                log.info("\(session.name, privacy: .public) reachable via \(String(describing: connection.endpoint), privacy: .public)")
            }
            connection.send(content: Data(reply), completion: .idempotent)

        case .dtls:
            guard let (id, session) = session(for: connection) else { return }
            let result = session.receiveDTLS(packet)
            for reply in result.replies { connection.send(content: Data(reply), completion: .idempotent) }
            if result.becameReady {
                log.info("\(session.name, privacy: .public) connected (DTLS-SRTP)")
                publish()
                onKeyFrameNeeded?()
            } else if case .failed(let reason) = session.dtlsState {
                log.error("DTLS with \(session.name, privacy: .public) failed: \(reason, privacy: .public)")
                sessions[id] = nil
                connections[id] = nil
                publish()
            }

        case .rtcp:
            if let (_, session) = session(for: connection), session.receiveRTCP(packet) { onKeyFrameNeeded?() }

        case .other:
            break
        }
    }

    private func session(for connection: NWConnection) -> (UUID, PeerSession)? {
        guard let id = connections.first(where: { $0.value === connection })?.key, let session = sessions[id] else { return nil }
        return (id, session)
    }

    private func startTimer() {
        let timer = DispatchSource.makeTimerSource(queue: queue)
        timer.schedule(deadline: .now() + 1, repeating: 1)
        timer.setEventHandler { [weak self] in self?.tick() }
        timer.resume()
        self.timer = timer
    }

    private func tick() {
        let now = Date()
        var changed = false
        for (id, session) in sessions {
            if session.isExpired(now: now) {
                log.info("\(session.name, privacy: .public) gone (no consent)")
                sessions[id] = nil
                connections.removeValue(forKey: id)?.cancel()
                changed = true
                continue
            }
            if let connection = connections[id] {
                for report in session.senderReports(now: now) { connection.send(content: Data(report), completion: .idempotent) }
            }
        }
        if changed { publish() }
    }

    private func publish() {
        let infos = sessions.values
            .sorted { $0.created < $1.created }
            .map { SessionInfo(id: $0.id, name: $0.name, role: $0.role, connected: $0.isMediaReady) }
        onSessionsChanged?(infos)
    }
}

final class OnceFlag: @unchecked Sendable {
    private let lock = NSLock()
    private var done = false
    func claim() -> Bool { lock.withLock { defer { done = true }; return !done } }
}