implement machine translation; various fixes and refactorings

2 files changed, 117 insertions, 0 deletions

diff --git a/src/translation/deepl.rs b/src/translation/deepl.rs
new file mode 100644
index 0000000..f2e84d7
--- /dev/null
+++ b/src/translation/deepl.rs
@@ -0,0 +1,106 @@
+use std::{collections::BTreeMap, time::Duration};
+
+use deepl::DeepLApi;
+use relm4::prelude::*;
+
+use crate::{
+    settings::Settings,
+    subtitles::{SUBTITLE_TRACKS, StreamIndex},
+    translation::TRANSLATIONS,
+};
+
+pub struct DeeplTranslator {
+    stream_ix: Option<StreamIndex>,
+    next_cues_to_translate: BTreeMap<StreamIndex, usize>,
+}
+
+#[derive(Debug)]
+pub enum DeeplTranslatorMsg {
+    SelectTrack(Option<StreamIndex>),
+    // this is only used to drive the async translation
+    DoTranslate,
+}
+
+impl AsyncComponent for DeeplTranslator {
+    type Init = ();
+    type Input = DeeplTranslatorMsg;
+    type Output = ();
+    type CommandOutput = ();
+    type Root = ();
+    type Widgets = ();
+
+    async fn init(
+        _init: Self::Init,
+        _root: Self::Root,
+        sender: relm4::AsyncComponentSender<Self>,
+    ) -> AsyncComponentParts<Self> {
+        let model = Self {
+            stream_ix: None,
+            next_cues_to_translate: BTreeMap::new(),
+        };
+
+        sender.input(DeeplTranslatorMsg::DoTranslate);
+
+        AsyncComponentParts { model, widgets: () }
+    }
+
+    async fn update(
+        &mut self,
+        message: Self::Input,
+        sender: AsyncComponentSender<Self>,
+        _root: &Self::Root,
+    ) {
+        match message {
+            DeeplTranslatorMsg::SelectTrack(stream_ix) => {
+                self.stream_ix = stream_ix;
+            }
+            DeeplTranslatorMsg::DoTranslate => self.do_translate(sender).await,
+        }
+    }
+
+    fn init_root() -> Self::Root {
+        ()
+    }
+}
+
+impl DeeplTranslator {
+    async fn do_translate(&mut self, sender: AsyncComponentSender<Self>) {
+        if let Some(stream_ix) = self.stream_ix {
+            let deepl = DeepLApi::with(&Settings::default().deepl_api_key()).new();
+
+            let next_cue_to_translate = self.next_cues_to_translate.entry(stream_ix).or_insert(0);
+
+            if let Some(cue) = {
+                SUBTITLE_TRACKS
+                    .read()
+                    .get(&stream_ix)
+                    .unwrap()
+                    .texts
+                    .get(*next_cue_to_translate)
+                    .cloned()
+            } {
+                match deepl
+                    .translate_text(cue, deepl::Lang::EN)
+                    .model_type(deepl::ModelType::PreferQualityOptimized)
+                    .await
+                {
+                    Ok(mut resp) => {
+                        TRANSLATIONS
+                            .write()
+                            .entry(stream_ix)
+                            .or_insert(Vec::new())
+                            .push(resp.translations.pop().unwrap().text);
+
+                        *next_cue_to_translate = *next_cue_to_translate + 1;
+                    }
+                    Err(e) => {
+                        log::error!("error fetching translation: {}", e)
+                    }
+                };
+            }
+        }
+
+        relm4::tokio::time::sleep(Duration::from_secs(1)).await;
+        sender.input(DeeplTranslatorMsg::DoTranslate);
+    }
+}
diff --git a/src/translation/mod.rs b/src/translation/mod.rs
new file mode 100644
index 0000000..4a1b358
--- /dev/null
+++ b/src/translation/mod.rs
@@ -0,0 +1,11 @@
+use std::collections::BTreeMap;
+
+use relm4::SharedState;
+
+use crate::subtitles::StreamIndex;
+
+pub mod deepl;
+
+pub use deepl::DeeplTranslator;
+
+pub static TRANSLATIONS: SharedState<BTreeMap<StreamIndex, Vec<String>>> = SharedState::new();