1 files changed, 72 insertions, 19 deletions

diff --git a/src/app.rs b/src/app.rs
index 49efd49..35a501e 100644
--- a/src/app.rs
+++ b/src/app.rs
@@ -1,4 +1,5 @@
 use adw::prelude::*;
+use cached::{DiskCache, IOCached};
 use relm4::{WorkerController, prelude::*};
 
 use crate::{
@@ -13,13 +14,15 @@ use crate::{
     subtitle_view::{SubtitleView, SubtitleViewMsg, SubtitleViewOutput},
     subtitles::{
         MetadataCollection, SUBTITLE_TRACKS, StreamIndex, SubtitleCue, SubtitleTrack,
-        TrackMetadata,
-        extraction::{SubtitleExtractor, SubtitleExtractorMsg, SubtitleExtractorOutput},
+        SubtitleTrackCollection, TrackMetadata,
+        extraction::{
+            ExtractionArgs, SubtitleExtractor, SubtitleExtractorMsg, SubtitleExtractorOutput,
+        },
         state::SubtitleState,
     },
     transcript::{Transcript, TranscriptMsg, TranscriptOutput},
-    translation::{DeeplTranslator, deepl::DeeplTranslatorMsg},
-    util::Tracker,
+    translation::{DeeplTranslator, TRANSLATIONS, deepl::DeeplTranslatorMsg},
+    util::{self, Tracker},
 };
 
 pub struct App {
@@ -34,6 +37,9 @@ pub struct App {
     open_url_dialog: Controller<OpenDialog>,
     subtitle_selection_dialog: Option<Controller<SubtitleSelectionDialog>>,
 
+    subtitle_extraction_args: Option<ExtractionArgs>,
+    subtitle_cache: DiskCache<ExtractionArgs, SubtitleTrackCollection>,
+
     primary_subtitle_state: SubtitleState,
     secondary_subtitle_state: SubtitleState,
 
@@ -153,6 +159,8 @@ impl SimpleComponent for App {
             },
         );
 
+        let subtitle_cache = util::make_cache("subtitles");
+
         let model = Self {
             root: root.clone(),
             player,
@@ -165,6 +173,9 @@ impl SimpleComponent for App {
             open_url_dialog,
             subtitle_selection_dialog: None,
 
+            subtitle_extraction_args: None,
+            subtitle_cache,
+
             primary_subtitle_state: SubtitleState::default(),
             secondary_subtitle_state: SubtitleState::default(),
 
@@ -193,6 +204,14 @@ impl SimpleComponent for App {
             }
             AppMsg::SubtitleExtractionComplete => {
                 log::info!("Subtitle extraction complete");
+                if let Some(ref args) = self.subtitle_extraction_args {
+                    if let Err(e) = self
+                        .subtitle_cache
+                        .cache_set(args.clone(), SUBTITLE_TRACKS.read().clone())
+                    {
+                        log::error!("error caching extracted subtitles: {}", e);
+                    }
+                }
             }
             AppMsg::ApplySubtitleSettings(settings) => {
                 self.primary_subtitle_state
@@ -249,6 +268,8 @@ impl SimpleComponent for App {
                 mut metadata,
                 whisper_stream_index,
             } => {
+                self.reset();
+
                 if let Some(ix) = whisper_stream_index {
                     let audio_metadata = metadata.audio.get(&ix).unwrap();
                     let subs_metadata = TrackMetadata {
@@ -263,17 +284,33 @@ impl SimpleComponent for App {
                     metadata.subtitles.insert(ix, subs_metadata);
                 }
 
-                self.player
-                    .sender()
-                    .send(PlayerMsg::SetUrl(url.clone()))
-                    .unwrap();
-                self.extractor
-                    .sender()
-                    .send(SubtitleExtractorMsg::ExtractFromUrl {
-                        url,
-                        whisper_stream_index,
-                    })
-                    .unwrap();
+                let extraction_args = ExtractionArgs {
+                    url: url.clone(),
+                    whisper_stream_index,
+                };
+
+                match self.subtitle_cache.cache_get(&extraction_args) {
+                    Ok(Some(track_collection)) => {
+                        log::debug!("subtitle cache hit");
+                        *(SUBTITLE_TRACKS.write()) = track_collection;
+                    }
+                    Ok(None) => {
+                        log::debug!("subtitle cache miss");
+                        self.extractor
+                            .sender()
+                            .send(SubtitleExtractorMsg::Extract(extraction_args.clone()))
+                            .unwrap();
+                    }
+                    Err(e) => {
+                        log::error!("error querying subtitle cache: {}", e);
+                        self.extractor
+                            .sender()
+                            .send(SubtitleExtractorMsg::Extract(extraction_args.clone()))
+                            .unwrap();
+                    }
+                }
+
+                self.subtitle_extraction_args = Some(extraction_args);
 
                 let subtitle_selection_dialog = SubtitleSelectionDialog::builder()
                     .launch((self.root.clone().into(), metadata))
@@ -283,12 +320,28 @@ impl SimpleComponent for App {
                         }
                     });
                 self.subtitle_selection_dialog = Some(subtitle_selection_dialog);
+
+                self.player.sender().send(PlayerMsg::SetUrl(url)).unwrap();
             }
         }
     }
 }
 
 impl App {
+    fn reset(&mut self) {
+        SUBTITLE_TRACKS.write().clear();
+        TRANSLATIONS.write().clear();
+
+        self.subtitle_selection_dialog = None;
+        self.subtitle_extraction_args = None;
+        self.primary_subtitle_state = SubtitleState::default();
+        self.secondary_subtitle_state = SubtitleState::default();
+        self.autopaused = false;
+        self.hovering_primary_cue = false;
+
+        // TODO also clear transcript?
+    }
+
     fn update_subtitle_states(&mut self, position: gst::ClockTime) {
         self.update_primary_subtitle_state(position);
         self.update_secondary_subtitle_state(position);
@@ -356,10 +409,10 @@ impl App {
 fn update_subtitle_state(state: &mut SubtitleState, position: gst::ClockTime) {
     if let Some(stream_ix) = state.stream_ix {
         let lock = SUBTITLE_TRACKS.read();
-        let track = lock.get(&stream_ix).unwrap();
-
-        update_last_time_ix(&track.start_times, &mut state.last_started_cue_ix, position);
-        update_last_time_ix(&track.end_times, &mut state.last_ended_cue_ix, position);
+        if let Some(track) = lock.get(&stream_ix) {
+            update_last_time_ix(&track.start_times, &mut state.last_started_cue_ix, position);
+            update_last_time_ix(&track.end_times, &mut state.last_ended_cue_ix, position);
+        }
     }
 }