Musubi/Packages/MusubiKit/Sources/NDIKit/NDIReceiver.swift

import CNDI
import CoreVideo
import Foundation
import MusubiCore
import OSLog

/// Receives video and/or audio from one NDI source on dedicated threads.
///
/// Video is handed out as IOSurface-backed pixel buffers, audio as planar float chunks.
/// Reconnects on its own when frames stop arriving (see `ConnectionWatchdog`).
public final class NDIReceiver: @unchecked Sendable {
    public enum Bandwidth: Sendable {
        case highest
        /// Audio and metadata only; the sender transmits no video to this receiver.
        case audioOnly
    }

    public struct Snapshot: Sendable, Equatable {
        public var state: ConnectionWatchdog.State
        public var stats: FrameStats
        public var width: Int
        public var height: Int
        public var frameRate: Double
        public var audioSampleRate: Int
        public var audioChannels: Int
        public var audioChunks: Int64
    }

    /// One block of planar 32-bit float audio. Only valid during the handler call.
    public struct AudioChunk: @unchecked Sendable {
        public let data: UnsafePointer<Float>
        /// Distance between channel planes, in samples.
        public let channelStride: Int
        public let channels: Int
        public let frames: Int
        public let sampleRate: Double
    }

    public typealias FrameHandler = @Sendable (CVPixelBuffer) -> Void
    public typealias AudioHandler = @Sendable (AudioChunk) -> Void

    public let source: Source
    private let instance: NDIlib_recv_instance_t
    private let onFrame: FrameHandler?
    private let onAudio: AudioHandler?
    private let log = Logger(subsystem: "studio.zielke.musubi", category: "ndi-recv")

    // Shared with the UI, guarded by `lock`.
    private let lock = NSLock()
    private var running = true
    private var watchdog: ConnectionWatchdog
    private var current: Snapshot

    // Only touched on the receive thread.
    private var pool: CVPixelBufferPool?
    private var poolKey: (width: Int, height: Int, format: OSType)?

    public init?(
        source: Source,
        receiverName: String,
        bandwidth: Bandwidth = .highest,
        onFrame: FrameHandler? = nil,
        onAudio: AudioHandler? = nil
    ) {
        let created: NDIlib_recv_instance_t? = source.ndiName.withCString { name in
            receiverName.withCString { recvName in
                var ndiSource = NDIlib_source_t()
                ndiSource.p_ndi_name = name
                var settings = NDIlib_recv_create_v3_t(
                    source_to_connect_to: ndiSource,
                    color_format: NDIlib_recv_color_format_UYVY_BGRA,
                    bandwidth: bandwidth == .audioOnly ? NDIlib_recv_bandwidth_audio_only : NDIlib_recv_bandwidth_highest,
                    allow_video_fields: false,
                    p_ndi_recv_name: recvName
                )
                return NDIlib_recv_create_v3(&settings)
            }
        }
        guard let created else { return nil }

        let now = ProcessInfo.processInfo.systemUptime
        self.source = source
        self.instance = created
        self.onFrame = onFrame
        self.onAudio = onAudio
        self.watchdog = ConnectionWatchdog(now: now)
        self.current = Snapshot(state: .connecting, stats: FrameStats(), width: 0, height: 0, frameRate: 0,
                                audioSampleRate: 0, audioChannels: 0, audioChunks: 0)

        // The first started loop also runs the watchdog. NDI allows capturing video and audio on separate threads.
        let videoPrimary = onFrame != nil
        if videoPrimary {
            startThread(name: "studio.zielke.musubi.ndi-recv") { [self] in run(capturesVideo: true, isPrimary: true) }
        }
        if onAudio != nil {
            startThread(name: "studio.zielke.musubi.ndi-audio") { [self] in run(capturesVideo: false, isPrimary: !videoPrimary) }
        }
    }

    private func startThread(name: String, _ body: @escaping @Sendable () -> Void) {
        let thread = Thread(block: body)
        thread.name = name
        thread.qualityOfService = .userInteractive
        thread.start()
    }

    deinit {
        NDIlib_recv_destroy(instance)
    }

    /// Stops the receive thread. The instance is destroyed once the last reference goes away.
    public func stop() {
        lock.withLock { running = false }
    }

    public var snapshot: Snapshot {
        lock.withLock {
            var snapshot = current
            snapshot.state = watchdog.state(at: ProcessInfo.processInfo.systemUptime)
            return snapshot
        }
    }

    // MARK: - Receive thread

    private func run(capturesVideo: Bool, isPrimary: Bool) {
        var lastPerformanceCheck = 0.0
        while lock.withLock({ running }) {
            let now: Double
            if capturesVideo {
                now = captureVideo()
            } else {
                now = captureAudio()
            }
            guard isPrimary else { continue }

            if now - lastPerformanceCheck >= 1 {
                lastPerformanceCheck = now
                var dropped = NDIlib_recv_performance_t()
                NDIlib_recv_get_performance(instance, nil, &dropped)
                let connected = NDIlib_recv_get_no_connections(instance) > 0
                lock.withLock {
                    current.stats.updateDroppedByReceiver(dropped.video_frames)
                    if connected { watchdog.connectionConfirmed(at: now) }
                }
            }

            if lock.withLock({ watchdog.shouldRetry(at: now) }) {
                log.notice("No frames from \(self.source.ndiName, privacy: .public), reconnecting")
                reconnect()
                lock.withLock { current.stats.resetTimeline() }
            }
        }
    }

    /// Captures one video frame (or times out). Returns the capture time.
    private func captureVideo() -> Double {
        var frame = NDIlib_video_frame_v2_t()
        let type = NDIlib_recv_capture_v3(instance, &frame, nil, nil, 250)
        let now = ProcessInfo.processInfo.systemUptime
        guard type == NDIlib_frame_type_video else { return now }

        if let onFrame, let buffer = makePixelBuffer(from: frame) {
            onFrame(buffer)
        }
        let rate = frame.frame_rate_D > 0 ? Double(frame.frame_rate_N) / Double(frame.frame_rate_D) : 0
        lock.withLock {
            watchdog.frameArrived(at: now)
            current.stats.record(timestamp: frame.timestamp, frameRateN: frame.frame_rate_N, frameRateD: frame.frame_rate_D)
            current.width = Int(frame.xres)
            current.height = Int(frame.yres)
            current.frameRate = rate
        }
        NDIlib_recv_free_video_v2(instance, &frame)
        return now
    }

    /// Captures one audio frame (or times out). Returns the capture time.
    private func captureAudio() -> Double {
        var frame = NDIlib_audio_frame_v3_t()
        let type = NDIlib_recv_capture_v3(instance, nil, &frame, nil, 250)
        let now = ProcessInfo.processInfo.systemUptime
        guard type == NDIlib_frame_type_audio else { return now }

        if frame.FourCC == NDIlib_FourCC_audio_type_FLTP, let data = frame.p_data, frame.no_samples > 0, frame.no_channels > 0 {
            data.withMemoryRebound(to: Float.self, capacity: Int(frame.channel_stride_in_bytes) / 4 * Int(frame.no_channels)) { samples in
                onAudio?(AudioChunk(
                    data: samples,
                    channelStride: Int(frame.channel_stride_in_bytes) / MemoryLayout<Float>.size,
                    channels: Int(frame.no_channels),
                    frames: Int(frame.no_samples),
                    sampleRate: Double(frame.sample_rate)
                ))
            }
        }
        lock.withLock {
            if onFrame == nil { watchdog.frameArrived(at: now) }
            current.audioSampleRate = Int(frame.sample_rate)
            current.audioChannels = Int(frame.no_channels)
            current.audioChunks += 1
        }
        NDIlib_recv_free_audio_v3(instance, &frame)
        return now
    }

    private func reconnect() {
        source.ndiName.withCString { name in
            var ndiSource = NDIlib_source_t()
            ndiSource.p_ndi_name = name
            NDIlib_recv_connect(instance, nil)
            NDIlib_recv_connect(instance, &ndiSource)
        }
    }

    private func makePixelBuffer(from frame: NDIlib_video_frame_v2_t) -> CVPixelBuffer? {
        let format: OSType
        let bytesPerPixel: Int
        switch frame.FourCC {
        case NDIlib_FourCC_video_type_UYVY:
            format = kCVPixelFormatType_422YpCbCr8 // '2vuy' has the same byte order as UYVY
            bytesPerPixel = 2
        case NDIlib_FourCC_video_type_BGRA, NDIlib_FourCC_video_type_BGRX:
            format = kCVPixelFormatType_32BGRA
            bytesPerPixel = 4
        default:
            return nil
        }
        let width = Int(frame.xres), height = Int(frame.yres)
        guard width > 0, height > 0, let source = frame.p_data,
              let pool = pixelBufferPool(width: width, height: height, format: format)
        else { return nil }

        var buffer: CVPixelBuffer?
        guard CVPixelBufferPoolCreatePixelBuffer(nil, pool, &buffer) == kCVReturnSuccess, let buffer else { return nil }

        CVPixelBufferLockBaseAddress(buffer, [])
        defer { CVPixelBufferUnlockBaseAddress(buffer, []) }
        guard let destination = CVPixelBufferGetBaseAddress(buffer) else { return nil }

        let sourceStride = Int(frame.line_stride_in_bytes)
        let destinationStride = CVPixelBufferGetBytesPerRow(buffer)
        if sourceStride == destinationStride {
            memcpy(destination, source, sourceStride * height)
        } else {
            let rowBytes = min(width * bytesPerPixel, sourceStride, destinationStride)
            for row in 0..<height {
                memcpy(destination + row * destinationStride, source + row * sourceStride, rowBytes)
            }
        }

        if format == kCVPixelFormatType_422YpCbCr8 {
            // NDI uses BT.709 for HD and BT.601 for SD.
            let hd = height >= 720
            CVBufferSetAttachment(buffer, kCVImageBufferYCbCrMatrixKey,
                                  hd ? kCVImageBufferYCbCrMatrix_ITU_R_709_2 : kCVImageBufferYCbCrMatrix_ITU_R_601_4, .shouldPropagate)
            CVBufferSetAttachment(buffer, kCVImageBufferColorPrimariesKey,
                                  hd ? kCVImageBufferColorPrimaries_ITU_R_709_2 : kCVImageBufferColorPrimaries_SMPTE_C, .shouldPropagate)
            CVBufferSetAttachment(buffer, kCVImageBufferTransferFunctionKey, kCVImageBufferTransferFunction_ITU_R_709_2, .shouldPropagate)
        }
        return buffer
    }

    private func pixelBufferPool(width: Int, height: Int, format: OSType) -> CVPixelBufferPool? {
        if let key = poolKey, key.width == width, key.height == height, key.format == format, let pool { return pool }
        let attributes: [CFString: Any] = [
            kCVPixelBufferWidthKey: width,
            kCVPixelBufferHeightKey: height,
            kCVPixelBufferPixelFormatTypeKey: format,
            kCVPixelBufferIOSurfacePropertiesKey: [CFString: Any](),
            kCVPixelBufferMetalCompatibilityKey: true,
        ]
        var newPool: CVPixelBufferPool?
        CVPixelBufferPoolCreate(nil, [kCVPixelBufferPoolMinimumBufferCountKey: 4] as CFDictionary, attributes as CFDictionary, &newPool)
        pool = newPool
        poolKey = (width, height, format)
        log.info("\(self.source.ndiName, privacy: .public): \(width)x\(height) \(format == kCVPixelFormatType_32BGRA ? "BGRA" : "UYVY", privacy: .public)")
        return newPool
    }
}