From 54327995e37cdda888a4dcb930736a681584d38b Mon Sep 17 00:00:00 2001
From: Jhen <developer@jhen.me>
Date: Sat, 30 Sep 2023 10:14:09 +0800
Subject: [PATCH] feat(ios, ts): add onNewSegments prop

---
 example/src/App.tsx    |  3 +++
 ios/RNWhisper.mm       |  9 ++++++++-
 src/NativeRNWhisper.ts |  2 +-
 src/index.ts           | 46 ++++++++++++++++++++++++++++++++++++++++--
 4 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/example/src/App.tsx b/example/src/App.tsx
index 4541c16..205a718 100644
--- a/example/src/App.tsx
+++ b/example/src/App.tsx
@@ -221,6 +221,9 @@ export default function App() {
                 onProgress: (cur) => {
                   log(`Transcribing progress: ${cur}%`)
                 },
+                // onNewSegments: (segments) => {
+                //   console.log('New segments:', segments)
+                // },
               })
               setStopTranscribe({ stop })
               const { result, segments } = await promise
diff --git a/ios/RNWhisper.mm b/ios/RNWhisper.mm
index 9343c06..4401f4b 100644
--- a/ios/RNWhisper.mm
+++ b/ios/RNWhisper.mm
@@ -87,6 +87,7 @@ - (NSDictionary *)constantsToExport
 - (NSArray *)supportedEvents {
   return@[
     @"@RNWhisper_onTranscribeProgress",
+    @"@RNWhisper_onTranscribeNewSegments",
     @"@RNWhisper_onRealtimeTranscribe",
     @"@RNWhisper_onRealtimeTranscribeEnd",
   ];
@@ -149,7 +150,13 @@ - (NSArray *)supportedEvents {
                 return;
             }
             dispatch_async(dispatch_get_main_queue(), ^{
-                // TODO
+                [self sendEventWithName:@"@RNWhisper_onTranscribeNewSegments"
+                    body:@{
+                        @"contextId": [NSNumber numberWithInt:contextId],
+                        @"jobId": [NSNumber numberWithInt:jobId],
+                        @"result": result
+                    }
+                ];
             });
         }
         onEnd: ^(int code) {
diff --git a/src/NativeRNWhisper.ts b/src/NativeRNWhisper.ts
index b7992a1..80fa1ea 100644
--- a/src/NativeRNWhisper.ts
+++ b/src/NativeRNWhisper.ts
@@ -68,7 +68,7 @@ export interface Spec extends TurboModule {
     contextId: number,
     jobId: number,
     path: string,
-    options: {}, // TranscribeOptions & { onProgress?: boolean }
+    options: {}, // TranscribeOptions & { onProgress?: boolean, onNewSegments?: boolean }
   ): Promise<TranscribeResult>;
   startRealtimeTranscribe(
     contextId: number,
diff --git a/src/index.ts b/src/index.ts
index 40bdb5c..70b23e2 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -26,16 +26,34 @@ export type { TranscribeOptions, TranscribeResult }
 
 
 const EVENT_ON_TRANSCRIBE_PROGRESS = '@RNWhisper_onTranscribeProgress'
+const EVENT_ON_TRANSCRIBE_NEW_SEGMENTS = '@RNWhisper_onTranscribeNewSegments'
 
 const EVENT_ON_REALTIME_TRANSCRIBE = '@RNWhisper_onRealtimeTranscribe'
 const EVENT_ON_REALTIME_TRANSCRIBE_END = '@RNWhisper_onRealtimeTranscribeEnd'
 
+export type TranscribeNewSegmentsResult = {
+  nNew: number
+  totalNNew: number
+  result: string
+  segments: TranscribeResult['segments']
+}
+
+export type TranscribeNewSegmentsNativeEvent = {
+  contextId: number
+  jobId: number
+  result: TranscribeNewSegmentsResult
+}
+
 // Fn -> Boolean in TranscribeFileNativeOptions
 export type TranscribeFileOptions = TranscribeOptions & {
   /**
    * Progress callback, the progress is between 0 and 100
    */
   onProgress?: (progress: number) => void
+  /**
+   * Callback when new segments are transcribed
+   */
+  onNewSegments?: (result: TranscribeNewSegmentsResult) => void
 }
 
 export type TranscribeProgressNativeEvent = {
@@ -155,7 +173,8 @@ export class WhisperContext {
     if (path.startsWith('file://')) path = path.slice(7)
     const jobId: number = Math.floor(Math.random() * 10000)
 
-    const { onProgress, ...rest } = options
+    const { onProgress, onNewSegments, ...rest } = options
+
     let progressListener: any
     let lastProgress: number = 0
     if (onProgress) {
@@ -175,16 +194,38 @@ export class WhisperContext {
         progressListener = null
       }
     }
+
+    let newSegmentsListener: any
+    if (onNewSegments) {
+      newSegmentsListener = EventEmitter.addListener(
+        EVENT_ON_TRANSCRIBE_NEW_SEGMENTS,
+        (evt: TranscribeNewSegmentsNativeEvent) => {
+          const { contextId, result } = evt
+          if (contextId !== this.id || evt.jobId !== jobId) return
+          onNewSegments(result)
+        },
+      )
+    }
+    const removeNewSegmenetsListener = () => {
+      if (newSegmentsListener) {
+        newSegmentsListener.remove()
+        newSegmentsListener = null
+      }
+    }
+
     return {
       stop: async () => {
         await RNWhisper.abortTranscribe(this.id, jobId)
         removeProgressListener()
+        removeNewSegmenetsListener()
       },
       promise: RNWhisper.transcribeFile(this.id, jobId, path, {
         ...rest,
-        onProgress: !!onProgress
+        onProgress: !!onProgress,
+        onNewSegments: !!onNewSegments
       }).then((result) => {
         removeProgressListener()
+        removeNewSegmenetsListener()
         if (!result.isAborted && lastProgress !== 100) {
           // Handle the case that the last progress event is not triggered
           onProgress?.(100)
@@ -192,6 +233,7 @@ export class WhisperContext {
         return result
       }).catch((e) => {
         removeProgressListener()
+        removeNewSegmenetsListener()
         throw e
       }),
     }