Musubi/Packages/MusubiKit/Sources/Intercom/Media/PreviewEncoder.swift

import CoreMedia
import QuartzCore
import CoreVideo
import Foundation
import OSLog
import VideoToolbox

/// Encodes the PGM preview for phones: scaled down, 25 fps, H.264 constrained baseline, no B-frames.
/// `encode` may be called from the NDI receive thread; work happens on an own queue.
final class PreviewEncoder: @unchecked Sendable {
    /// Called with the NAL units of one access unit and its 90 kHz timestamp.
    var onAccessUnit: (@Sendable ([[UInt8]], UInt32) -> Void)?

    let width: Int32
    let height: Int32
    let frameRate: Double
    private let bitRate: Int
    private let queue = DispatchQueue(label: "studio.zielke.musubi.intercom.video", qos: .userInitiated)
    private let log = Logger(subsystem: "studio.zielke.musubi", category: "intercom")

    // queue only
    private var session: VTCompressionSession?
    private var transfer: VTPixelTransferSession?
    private var pool: CVPixelBufferPool?
    private var forceKeyFrame = true
    private var lastFrameTime: Double = -1
    private var busy = false
    private let start = CACurrentMediaTime()

    init(width: Int32 = 960, height: Int32 = 540, frameRate: Double = 25, bitRate: Int = 1_200_000) {
        self.width = width
        self.height = height
        self.frameRate = frameRate
        self.bitRate = bitRate
    }

    deinit {
        if let session { VTCompressionSessionInvalidate(session) }
        if let transfer { VTPixelTransferSessionInvalidate(transfer) }
    }

    func requestKeyFrame() {
        queue.async { self.forceKeyFrame = true }
    }

    /// Offers a frame; frames above the target rate, or while the encoder is still busy, are dropped.
    func encode(_ pixelBuffer: CVPixelBuffer) {
        let now = CACurrentMediaTime() - start
        // CVPixelBuffer is immutable once handed over by the receiver.
        nonisolated(unsafe) let pixelBuffer = pixelBuffer
        queue.async {
            guard !self.busy, now - self.lastFrameTime >= 1 / self.frameRate - 0.004 else { return }
            self.lastFrameTime = now
            self.encodeOnQueue(pixelBuffer, time: now)
        }
    }

    private func encodeOnQueue(_ source: CVPixelBuffer, time: Double) {
        guard let session = session ?? makeSession(), let scaled = scale(source) else { return }
        var options: CFDictionary?
        if forceKeyFrame {
            options = [kVTEncodeFrameOptionKey_ForceKeyFrame: true] as CFDictionary
            forceKeyFrame = false
        }
        busy = true
        let pts = CMTime(seconds: time, preferredTimescale: 90_000)
        let status = VTCompressionSessionEncodeFrame(session, imageBuffer: scaled, presentationTimeStamp: pts, duration: .invalid,
                                                     frameProperties: options, infoFlagsOut: nil) { [weak self] status, _, sample in
            guard let self else { return }
            self.queue.async { self.busy = false }
            guard status == noErr, let sample else { return }
            self.emit(sample, timestamp: UInt32(truncatingIfNeeded: Int64(time * 90_000)))
        }
        if status != noErr {
            busy = false
            log.error("Preview encode failed: \(status)")
        }
    }

    private func makeSession() -> VTCompressionSession? {
        var session: VTCompressionSession?
        let status = VTCompressionSessionCreate(allocator: nil, width: width, height: height, codecType: kCMVideoCodecType_H264,
                                                encoderSpecification: nil, imageBufferAttributes: nil, compressedDataAllocator: nil,
                                                outputCallback: nil, refcon: nil, compressionSessionOut: &session)
        guard status == noErr, let session else {
            log.error("Preview encoder unavailable: \(status)")
            return nil
        }
        let properties: [CFString: Any] = [
            kVTCompressionPropertyKey_RealTime: true,
            kVTCompressionPropertyKey_ProfileLevel: kVTProfileLevel_H264_ConstrainedBaseline_AutoLevel,
            kVTCompressionPropertyKey_AllowFrameReordering: false,
            kVTCompressionPropertyKey_AverageBitRate: bitRate,
            kVTCompressionPropertyKey_ExpectedFrameRate: frameRate,
            kVTCompressionPropertyKey_MaxKeyFrameInterval: Int(frameRate * 2),
            kVTCompressionPropertyKey_MaxKeyFrameIntervalDuration: 2,
        ]
        for (key, value) in properties { VTSessionSetProperty(session, key: key, value: value as CFTypeRef) }
        VTCompressionSessionPrepareToEncodeFrames(session)
        self.session = session
        return session
    }

    /// Scales any NDI frame (BGRA or UYVY, any size) into a pooled NV12 buffer of the preview size.
    private func scale(_ source: CVPixelBuffer) -> CVPixelBuffer? {
        if transfer == nil {
            var newTransfer: VTPixelTransferSession?
            VTPixelTransferSessionCreate(allocator: nil, pixelTransferSessionOut: &newTransfer)
            if let newTransfer {
                VTSessionSetProperty(newTransfer, key: kVTPixelTransferPropertyKey_ScalingMode, value: kVTScalingMode_Letterbox)
            }
            transfer = newTransfer
        }
        if pool == nil {
            let attributes: [CFString: Any] = [
                kCVPixelBufferPixelFormatTypeKey: kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange,
                kCVPixelBufferWidthKey: width,
                kCVPixelBufferHeightKey: height,
                kCVPixelBufferIOSurfacePropertiesKey: [:] as CFDictionary,
            ]
            CVPixelBufferPoolCreate(nil, nil, attributes as CFDictionary, &pool)
        }
        guard let transfer, let pool else { return nil }
        var target: CVPixelBuffer?
        CVPixelBufferPoolCreatePixelBuffer(nil, pool, &target)
        guard let target, VTPixelTransferSessionTransferImage(transfer, from: source, to: target) == noErr else { return nil }
        return target
    }

    private func emit(_ sample: CMSampleBuffer, timestamp: UInt32) {
        guard let block = CMSampleBufferGetDataBuffer(sample) else { return }
        var length = 0
        var pointer: UnsafeMutablePointer<CChar>?
        guard CMBlockBufferGetDataPointer(block, atOffset: 0, lengthAtOffsetOut: nil, totalLengthOut: &length,
                                          dataPointerOut: &pointer) == noErr, let pointer else { return }
        let avcc = [UInt8](UnsafeRawBufferPointer(start: pointer, count: length))
        var nalUnits = H264Packetizer.nalUnits(avcc: avcc)

        let attachments = CMSampleBufferGetSampleAttachmentsArray(sample, createIfNecessary: false) as? [[CFString: Any]]
        let isKeyFrame = !(attachments?.first?[kCMSampleAttachmentKey_NotSync] as? Bool ?? false)
        if isKeyFrame, let format = CMSampleBufferGetFormatDescription(sample) {
            nalUnits = Self.parameterSets(format) + nalUnits
        }
        onAccessUnit?(nalUnits, timestamp)
    }

    private static func parameterSets(_ format: CMFormatDescription) -> [[UInt8]] {
        var count = 0
        CMVideoFormatDescriptionGetH264ParameterSetAtIndex(format, parameterSetIndex: 0, parameterSetPointerOut: nil,
                                                           parameterSetSizeOut: nil, parameterSetCountOut: &count, nalUnitHeaderLengthOut: nil)
        return (0..<count).compactMap { index in
            var pointer: UnsafePointer<UInt8>?
            var size = 0
            guard CMVideoFormatDescriptionGetH264ParameterSetAtIndex(format, parameterSetIndex: index, parameterSetPointerOut: &pointer,
                                                                     parameterSetSizeOut: &size, parameterSetCountOut: nil,
                                                                     nalUnitHeaderLengthOut: nil) == noErr, let pointer else { return nil }
            return Array(UnsafeBufferPointer(start: pointer, count: size))
        }
    }
}