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

import CNDI
import CoreVideo
import Foundation
import MusubiCore
import OSLog

/// Receives video from one NDI source on a dedicated thread and hands out IOSurface-backed pixel buffers.
///
/// Reconnects on its own when frames stop arriving (see `ConnectionWatchdog`).
public final class NDIReceiver: @unchecked Sendable {
    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 typealias FrameHandler = @Sendable (CVPixelBuffer) -> Void

    public let source: Source
    private let instance: NDIlib_recv_instance_t
    private let onFrame: FrameHandler
    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, onFrame: @escaping FrameHandler) {
        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: 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.watchdog = ConnectionWatchdog(now: now)
        self.current = Snapshot(state: .connecting, stats: FrameStats(), width: 0, height: 0, frameRate: 0)

        let thread = Thread { [self] in run() }
        thread.name = "studio.zielke.musubi.ndi-recv"
        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() {
        var lastPerformanceCheck = 0.0
        while lock.withLock({ running }) {
            var frame = NDIlib_video_frame_v2_t()
            let type = NDIlib_recv_capture_v3(instance, &frame, nil, nil, 250)
            let now = ProcessInfo.processInfo.systemUptime

            if type == NDIlib_frame_type_video {
                if 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)
            }

            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() }
            }
        }
    }

    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
    }
}