From 8e46d793767c1f3bb80fbb4d9180ef3eb335b64c Mon Sep 17 00:00:00 2001
From: Sami Samhuri <sami@samhuri.net>
Date: Sat, 17 Aug 2024 15:39:53 -0700
Subject: [PATCH] Add support for writing metadata on the asset

---
 SJSAssetExportSession/ExportSession.swift     |  6 ++++
 SJSAssetExportSession/SampleWriter.swift      |  2 ++
 .../SJSAssetExportSessionTests.swift          | 36 +++++++++++++++++++
 3 files changed, 44 insertions(+)

diff --git a/SJSAssetExportSession/ExportSession.swift b/SJSAssetExportSession/ExportSession.swift
index 251f3de..508eefa 100644
--- a/SJSAssetExportSession/ExportSession.swift
+++ b/SJSAssetExportSession/ExportSession.swift
@@ -25,6 +25,7 @@ public final class ExportSession: @unchecked Sendable {
     public func export(
         asset: sending AVAsset,
         optimizeForNetworkUse: Bool = false,
+        metadata: sending [AVMetadataItem] = [],
         timeRange: CMTimeRange? = nil,
         audio: sending AudioOutputSettings = .default,
         video: sending VideoOutputSettings,
@@ -42,6 +43,7 @@ public final class ExportSession: @unchecked Sendable {
             videoComposition: videoComposition,
             timeRange: timeRange,
             optimizeForNetworkUse: optimizeForNetworkUse,
+            metadata: metadata,
             outputURL: outputURL,
             fileType: fileType
         )
@@ -61,6 +63,8 @@ public final class ExportSession: @unchecked Sendable {
 
        - optimizeForNetworkUse: Setting this value to `true` writes the output file in a form that enables a player to begin playing the media after downloading only a small portion of it. Defaults to `false`.
 
+       - metadata: Optional array of `AVMetadataItem`s to be written out with the exported asset.
+
        - timeRange: Providing a time range exports a subset of the asset instead of the entire duration, which is the default behaviour.
 
        - audioOutputSettings: Audio settings using [audio settings keys from AVFoundation](https://developer.apple.com/documentation/avfoundation/audio_settings) and values must be suitable for consumption by Objective-C. Required keys are:
@@ -83,6 +87,7 @@ public final class ExportSession: @unchecked Sendable {
     public func export(
         asset: sending AVAsset,
         optimizeForNetworkUse: Bool = false,
+        metadata: sending [AVMetadataItem] = [],
         timeRange: CMTimeRange? = nil,
         audioOutputSettings: [String: (any Sendable)],
         mix: sending AVAudioMix? = nil,
@@ -118,6 +123,7 @@ public final class ExportSession: @unchecked Sendable {
             videoComposition: videoComposition,
             timeRange: timeRange,
             optimizeForNetworkUse: optimizeForNetworkUse,
+            metadata: metadata,
             outputURL: outputURL,
             fileType: fileType
         )
diff --git a/SJSAssetExportSession/SampleWriter.swift b/SJSAssetExportSession/SampleWriter.swift
index 76bd36f..e9873f7 100644
--- a/SJSAssetExportSession/SampleWriter.swift
+++ b/SJSAssetExportSession/SampleWriter.swift
@@ -53,6 +53,7 @@ actor SampleWriter {
         videoComposition: sending AVVideoComposition,
         timeRange: CMTimeRange? = nil,
         optimizeForNetworkUse: Bool = false,
+        metadata: [AVMetadataItem] = [],
         outputURL: URL,
         fileType: AVFileType
     ) async throws {
@@ -66,6 +67,7 @@ actor SampleWriter {
         }
         let writer = try AVAssetWriter(outputURL: outputURL, fileType: fileType)
         writer.shouldOptimizeForNetworkUse = optimizeForNetworkUse
+        writer.metadata = metadata
 
         let audioTracks = try await asset.loadTracks(withMediaType: .audio)
         // Audio is optional so only validate output settings when it's applicable.
diff --git a/SJSAssetExportSessionTests/SJSAssetExportSessionTests.swift b/SJSAssetExportSessionTests/SJSAssetExportSessionTests.swift
index 19b463a..8a4c639 100644
--- a/SJSAssetExportSessionTests/SJSAssetExportSessionTests.swift
+++ b/SJSAssetExportSessionTests/SJSAssetExportSessionTests.swift
@@ -6,6 +6,7 @@
 //
 
 internal import AVFoundation
+import CoreLocation
 @testable import SJSAssetExportSession
 import Testing
 
@@ -321,4 +322,39 @@ final class ExportSessionTests {
         try? await task.value // Wait for task to complete
         NSLog("Task has finished executing")
     }
+
+    @Test func test_writing_metadata() async throws {
+        let sourceURL = resourceURL(named: "test-720p-h264-24fps", withExtension: "mov")
+        let destinationURL = makeTemporaryURL()
+        let locationMetadata = AVMutableMetadataItem()
+        locationMetadata.key = AVMetadataKey.commonKeyLocation.rawValue as NSString
+        locationMetadata.keySpace = .common
+        locationMetadata.value = "+48.50176+123.34368/" as NSString
+
+        let subject = ExportSession()
+        try await subject.export(
+            asset: makeAsset(url: sourceURL),
+            metadata: [locationMetadata],
+            video: .codec(.h264, size: CGSize(width: 1280, height: 720)),
+            to: destinationURL.url,
+            as: .mov
+        )
+
+        let exportedAsset = AVURLAsset(url: destinationURL.url)
+        let exportedMetadata = try await exportedAsset.load(.metadata)
+        print(exportedMetadata)
+        #expect(exportedMetadata.count == 1)
+        let metadataValue = try await exportedMetadata.first(where: { item in
+            item.key as! String == AVMetadataKey.quickTimeMetadataKeyLocationISO6709.rawValue
+        })?.load(.value) as? NSString
+        #expect(metadataValue == "+48.50176+123.34368/")
+
+        let exportedCommonMetadata = try await exportedAsset.load(.commonMetadata)
+        print(exportedCommonMetadata)
+        #expect(exportedCommonMetadata.count == 1)
+        let commonMetadataValue = try await exportedCommonMetadata.first(where: { item in
+            item.commonKey == .commonKeyLocation
+        })?.load(.value) as? NSString
+        #expect(commonMetadataValue == "+48.50176+123.34368/")
+    }
 }