import Foundation
import AVFoundation
import Metal
import CoreVideo
/// Real ClipRecorder: AVAssetWriter + pixel-buffer-pool + CVMetalTextureCache.
/// The caller renders each frame into the texture handed to `renderInto`.
final class VideoClipRecorder: ClipRecorder {
struct LivePhoto { let assetID: String; let keyFrameIndex: Int }
private let device: MTLDevice
private let url: URL
private let width: Int, height: Int, fps: Int, codec: VideoCodec
private let renderInto: (MTLTexture) throws -> Void
private let livePhoto: LivePhoto?
private var writer: AVAssetWriter!
private var input: AVAssetWriterInput!
private var adaptor: AVAssetWriterInputPixelBufferAdaptor!
private var textureCache: CVMetalTextureCache!
private var metaInput: AVAssetWriterInput?
private var metaAdaptor: AVAssetWriterInputMetadataAdaptor?
private(set) var keyFrameImage: CGImage?
init(device: MTLDevice, url: URL, width: Int, height: Int, fps: Int, codec: VideoCodec,
renderInto: @escaping (MTLTexture) throws -> Void,
livePhoto: LivePhoto? = nil) {
self.device = device; self.url = url
self.width = width; self.height = height; self.fps = fps; self.codec = codec
self.renderInto = renderInto
self.livePhoto = livePhoto
}
func begin() throws {
try? FileManager.default.removeItem(at: url)
let fileType: AVFileType = codec == .hevc ? .mov : .mp4
writer = try AVAssetWriter(outputURL: url, fileType: fileType)
let codecType: AVVideoCodecType = codec == .hevc ? .hevc : .h264
let compression: [String: Any] = [
AVVideoAverageBitRateKey: bitrate(),
AVVideoMaxKeyFrameIntervalKey: fps * 2,
]
let settings: [String: Any] = [
AVVideoCodecKey: codecType,
AVVideoWidthKey: width,
AVVideoHeightKey: height,
AVVideoCompressionPropertiesKey: compression,
AVVideoColorPropertiesKey: [
AVVideoColorPrimariesKey: AVVideoColorPrimaries_ITU_R_709_2,
AVVideoTransferFunctionKey: AVVideoTransferFunction_ITU_R_709_2,
AVVideoYCbCrMatrixKey: AVVideoYCbCrMatrix_ITU_R_709_2,
],
]
input = AVAssetWriterInput(mediaType: .video, outputSettings: settings)
input.expectsMediaDataInRealTime = false
let pbAttrs: [String: Any] = [
kCVPixelBufferPixelFormatTypeKey as String: kCVPixelFormatType_32BGRA,
kCVPixelBufferMetalCompatibilityKey as String: true,
kCVPixelBufferIOSurfacePropertiesKey as String: [:],
kCVPixelBufferWidthKey as String: width,
kCVPixelBufferHeightKey as String: height,
]
adaptor = AVAssetWriterInputPixelBufferAdaptor(assetWriterInput: input, sourcePixelBufferAttributes: pbAttrs)
guard writer.canAdd(input) else { throw CaptureError.writer("cannot add input") }
writer.add(input)
var cache: CVMetalTextureCache?
let cacheStatus = CVMetalTextureCacheCreate(nil, nil, device, nil, &cache)
guard cacheStatus == kCVReturnSuccess, let cache else {
throw CaptureError.writer("texture cache create failed (\(cacheStatus))")
}
textureCache = cache
if let lp = livePhoto {
writer.metadata = [contentIdentifierItem(lp.assetID)]
let spec: [String: Any] = [
kCMMetadataFormatDescriptionMetadataSpecificationKey_Identifier as String: "mdta/com.apple.quicktime.still-image-time",
kCMMetadataFormatDescriptionMetadataSpecificationKey_DataType as String: "com.apple.metadata.datatype.int8",
]
var desc: CMFormatDescription?
CMMetadataFormatDescriptionCreateWithMetadataSpecifications(
allocator: kCFAllocatorDefault,
metadataType: kCMMetadataFormatType_Boxed,
metadataSpecifications: [spec] as CFArray,
formatDescriptionOut: &desc)
let mInput = AVAssetWriterInput(mediaType: .metadata, outputSettings: nil, sourceFormatHint: desc)
mInput.expectsMediaDataInRealTime = false
guard writer.canAdd(mInput) else { throw CaptureError.writer("cannot add metadata input") }
writer.add(mInput)
metaInput = mInput
metaAdaptor = AVAssetWriterInputMetadataAdaptor(assetWriterInput: mInput)
}
guard writer.startWriting() else { throw CaptureError.writer("\(writer.error?.localizedDescription ?? "startWriting")") }
writer.startSession(atSourceTime: .zero)
if let lp = livePhoto, let adaptor = metaAdaptor {
let keyTime = CMTime(value: CMTimeValue(lp.keyFrameIndex), timescale: CMTimeScale(fps))
let oneFrame = CMTime(value: 1, timescale: CMTimeScale(fps))
adaptor.append(AVTimedMetadataGroup(
items: [stillImageTimeItem()],
timeRange: CMTimeRange(start: keyTime, duration: oneFrame)))
}
}
func appendFrame(index: Int, pts: CMTime) throws {
while !input.isReadyForMoreMediaData {
if writer.status == .failed { throw CaptureError.writer("\(writer.error?.localizedDescription ?? "failed")") }
Thread.sleep(forTimeInterval: 0.002)
}
try autoreleasepool {
guard let pool = adaptor.pixelBufferPool else { throw CaptureError.writer("no pool") }
var pbOut: CVPixelBuffer?
let pbStatus = CVPixelBufferPoolCreatePixelBuffer(nil, pool, &pbOut)
guard pbStatus == kCVReturnSuccess, let pb = pbOut else {
throw CaptureError.writer("pixel buffer pool failed (\(pbStatus))")
}
var cvTexOut: CVMetalTexture?
let r = CVMetalTextureCacheCreateTextureFromImage(
nil, textureCache, pb, nil, .bgra8Unorm, width, height, 0, &cvTexOut)
guard r == kCVReturnSuccess, let cvTex = cvTexOut,
let tex = CVMetalTextureGetTexture(cvTex) else { throw CaptureError.writer("texture from buffer") }
try renderInto(tex)
if let lp = livePhoto, index == lp.keyFrameIndex {
keyFrameImage = ImageWriter.cgImage(from: pb)
}
if !adaptor.append(pb, withPresentationTime: pts) {
throw CaptureError.writer("\(writer.error?.localizedDescription ?? "append")")
}
CVMetalTextureCacheFlush(textureCache, 0)
}
}
func finish() async throws -> URL {
metaInput?.markAsFinished()
input.markAsFinished()
await withCheckedContinuation { (c: CheckedContinuation<Void, Never>) in
writer.finishWriting { c.resume() }
}
if writer.status != .completed { throw CaptureError.writer("\(writer.error?.localizedDescription ?? "finish")") }
return url
}
private func bitrate() -> Int {
let pixels = width * height
let base = codec == .hevc ? 8.0 : 14.0
return Int(base * 1_000_000 * Double(pixels) / Double(1920 * 1080))
}
private func contentIdentifierItem(_ id: String) -> AVMetadataItem {
let i = AVMutableMetadataItem()
i.key = "com.apple.quicktime.content.identifier" as NSString
i.keySpace = AVMetadataKeySpace(rawValue: "mdta")
i.value = id as NSString
i.dataType = "com.apple.metadata.datatype.UTF-8"
return i
}
private func stillImageTimeItem() -> AVMetadataItem {
let i = AVMutableMetadataItem()
i.key = "com.apple.quicktime.still-image-time" as NSString
i.keySpace = AVMetadataKeySpace(rawValue: "mdta")
i.value = 0 as NSNumber
i.dataType = "com.apple.metadata.datatype.int8"
return i
}
}