summary refs log tree commit diff
path: root/src/translation/deepl.rs
blob: f2e84d793ef9d09a2f10c378cb814020c4c12a06 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
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);
    }
}