Add subtitle extractor (#22)

* subtitle extractor

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "movie_collection_rust"
-version = "0.10.32"
+version = "0.10.33"
 authors = ["Daniel Boline <[email protected]>"]
 edition = "2018"
 

movie_collection_http/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "movie_collection_http"
-version = "0.10.32"
+version = "0.10.33"
 authors = ["Daniel Boline <[email protected]>"]
 edition = "2018"
 

movie_collection_http/src/movie_queue_app.rs

@@ -30,16 +30,16 @@ use super::{
         find_new_episodes, frontpage, imdb_episodes_route, imdb_episodes_update,
         imdb_ratings_route, imdb_ratings_set_source, imdb_ratings_update, imdb_show,
         last_modified_route, movie_collection_route, movie_collection_update, movie_queue,
-        movie_queue_delete, movie_queue_play, movie_queue_remcom_directory_file,
-        movie_queue_remcom_file, movie_queue_route, movie_queue_show, movie_queue_transcode,
-        movie_queue_transcode_cleanup, movie_queue_transcode_directory, movie_queue_transcode_file,
-        movie_queue_transcode_status, movie_queue_transcode_status_file_list,
-        movie_queue_transcode_status_procs, movie_queue_update, music_collection,
-        music_collection_update, plex_detail, plex_events, plex_events_update, plex_filename,
-        plex_filename_update, plex_list, plex_metadata, plex_metadata_update, plex_webhook,
-        refresh_auth, scripts_js, trakt_auth_url, trakt_cal, trakt_callback, trakt_watched_action,
-        trakt_watched_list, trakt_watched_seasons, trakt_watchlist, trakt_watchlist_action,
-        tvshows, user,
+        movie_queue_delete, movie_queue_extract_subtitle, movie_queue_play,
+        movie_queue_remcom_directory_file, movie_queue_remcom_file, movie_queue_route,
+        movie_queue_show, movie_queue_transcode, movie_queue_transcode_cleanup,
+        movie_queue_transcode_directory, movie_queue_transcode_file, movie_queue_transcode_status,
+        movie_queue_transcode_status_file_list, movie_queue_transcode_status_procs,
+        movie_queue_update, music_collection, music_collection_update, plex_detail, plex_events,
+        plex_events_update, plex_filename, plex_filename_update, plex_list, plex_metadata,
+        plex_metadata_update, plex_webhook, refresh_auth, scripts_js, trakt_auth_url, trakt_cal,
+        trakt_callback, trakt_watched_action, trakt_watched_list, trakt_watched_seasons,
+        trakt_watchlist, trakt_watchlist_action, tvshows, user,
     },
 };
 
@@ -138,6 +138,7 @@ fn get_full_path(app: &AppState) -> BoxedFilter<(impl Reply,)> {
     let plex_metadata_update_path = plex_metadata_update(app.clone()).boxed();
     let music_collection_path = music_collection(app.clone()).boxed();
     let music_collection_update_path = music_collection_update(app.clone()).boxed();
+    let movie_queue_extract_subtitle_path = movie_queue_extract_subtitle(app.clone()).boxed();
 
     let list_path = frontpage_path
         .or(scripts_js_path)
@@ -166,7 +167,8 @@ fn get_full_path(app: &AppState) -> BoxedFilter<(impl Reply,)> {
         .or(plex_metadata_path)
         .or(plex_metadata_update_path)
         .or(music_collection_path)
-        .or(music_collection_update_path);
+        .or(music_collection_update_path)
+        .or(movie_queue_extract_subtitle_path);
     let auth_url_path = trakt_auth_url(app.clone()).boxed();
     let trakt_callback_path = trakt_callback(app.clone()).boxed();
     let refresh_auth_path = refresh_auth(app.clone()).boxed();

movie_collection_http/src/movie_queue_elements.rs

@@ -2081,11 +2081,14 @@ fn LocalFileElement(
         .iter()
         .enumerate()
         .map(|(idx, f)| {
-            let f_key = f
-                .replace(".mkv", "")
-                .replace(".m4v", "")
-                .replace(".avi", "")
-                .replace(".mp4", "");
+            let mut f_key = f.clone();
+            for suffix in [".mkv", ".m4v", ".avi", ".mp4"] {
+                if let Some(s) = f_key.strip_suffix(suffix) {
+                    f_key = s.into();
+                    break;
+                }
+            }
+            let mut subtitle_selector = None;
             let button = if file_map.contains_key(f_key.as_str()) {
                 rsx! {
                     button {
@@ -2134,6 +2137,38 @@ fn LocalFileElement(
                     None => rsx! {"unknown"},
                 }
             } else {
+                if let Some(subtitles) = file_lists.subtitles.get(f.as_str()) {
+                    let nlines = subtitles.iter().find_map(|(_, n)| *n);
+                    let options = subtitles.iter().enumerate().map(|(i, (s, _))| {
+                        let mkv_number = s.number;
+                        let label =
+                            format_sstr!("{} {} {} {}", s.number, s.language, s.codec_id, s.name);
+                        rsx! {
+                            option {
+                                key: "subtitle-key-{f}-{i}",
+                                value: "{mkv_number}",
+                                "{label}"
+                            }
+                        }
+                    });
+                    let title = if let Some(n) = nlines {
+                        format_sstr!("re-extract subtitles {n}")
+                    } else {
+                        "extract subtitles".into()
+                    };
+                    subtitle_selector.replace(rsx! {
+                        select {
+                            id: "subtitle-selector-{f}",
+                            {options},
+                        },
+                        button {
+                            "type": "submit",
+                            id: "subtitle-button-{f}",
+                            "onclick": "extract_subtitles('{f}')",
+                            "{title}",
+                        }
+                    });
+                }
                 rsx! {
                     button {
                         "type": "submit",
@@ -2149,6 +2184,7 @@ fn LocalFileElement(
                     key: "flist-key-{idx}",
                     td {"{f}"},
                     td {{button}},
+                    td {{subtitle_selector}},
                 }
             }
         });

movie_collection_http/src/movie_queue_routes.rs

@@ -29,6 +29,7 @@ use movie_collection_lib::{
     imdb_episodes::ImdbEpisodes,
     imdb_ratings::ImdbRatings,
     make_list::FileLists,
+    mkv_utils::MkvTrack,
     movie_collection::{LastModifiedResponse, MovieCollection, TvShowsResult},
     movie_queue::{MovieQueueDB, MovieQueueResult},
     music_collection::MusicCollection,
@@ -1869,3 +1870,42 @@ pub async fn music_collection_update(
     task.await.ok();
     Ok(HtmlBase::new("Success").into())
 }
+
+#[derive(RwebResponse)]
+#[response(
+    description = "Extract Subtitles",
+    content = "html",
+    status = "Created"
+)]
+struct ExtractSubtitlesResponse(HtmlBase<StackString, Error>);
+
+#[post("/list/transcode/subtitle/{filename}/{index}")]
+pub async fn movie_queue_extract_subtitle(
+    filename: StackString,
+    index: u64,
+    #[filter = "LoggedUser::filter"] user: LoggedUser,
+    #[data] state: AppState,
+) -> WarpResult<ExtractSubtitlesResponse> {
+    let task = user
+        .store_url_task(
+            state.trakt.get_client(),
+            &state.config,
+            &format_sstr!("/list/transcode/subtitle/{filename}/{index}"),
+        )
+        .await;
+
+    let input_path = state
+        .config
+        .home_dir
+        .join("Documents")
+        .join("movies")
+        .join(&filename);
+
+    let input_file: StackString = input_path.to_string_lossy().into();
+    let output = MkvTrack::extract_subtitles_from_mkv(&input_file, index)
+        .await
+        .map_err(Into::<Error>::into)?;
+    task.await.ok();
+
+    Ok(HtmlBase::new(output).into())
+}

movie_collection_lib/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "movie_collection_lib"
-version = "0.10.32"
+version = "0.10.33"
 authors = ["Daniel Boline <[email protected]>"]
 edition = "2018"
 

movie_collection_lib/src/lib.rs

@@ -15,6 +15,7 @@ pub mod imdb_ratings;
 pub mod imdb_utils;
 pub mod make_list;
 pub mod make_queue;
+pub mod mkv_utils;
 pub mod movie_collection;
 pub mod movie_queue;
 pub mod music_collection;

movie_collection_lib/src/make_list.rs

@@ -18,6 +18,7 @@ use tokio_stream::{wrappers::ReadDirStream, StreamExt};
 
 use crate::{
     config::Config,
+    mkv_utils::{MkvTrack, TrackType},
     movie_collection::MovieCollection,
     pgpool::PgPool,
     transcode_service::transcode_status,
@@ -28,6 +29,7 @@ use crate::{
 pub struct FileLists {
     pub local_file_list: Vec<StackString>,
     pub file_list: Vec<PathBuf>,
+    pub subtitles: HashMap<StackString, Vec<(MkvTrack, Option<usize>)>>,
 }
 
 impl FileLists {
@@ -71,10 +73,12 @@ impl FileLists {
         let patterns: Vec<_> = local_file_list
             .iter()
             .map(|f| {
-                f.replace(".mkv", "")
-                    .replace(".m4v", "")
-                    .replace(".avi", "")
-                    .replace(".mp4", "")
+                for suffix in [".mkv", ".m4v", ".avi", ".mp4"] {
+                    if let Some(s) = f.strip_suffix(suffix) {
+                        return s.into();
+                    }
+                }
+                f.clone()
             })
             .collect();
 
@@ -107,9 +111,33 @@ impl FileLists {
             file_list?
         };
 
+        let mut subtitles: HashMap<StackString, Vec<(MkvTrack, Option<usize>)>> = HashMap::new();
+
+        for f in &local_file_list {
+            if f.ends_with(".mkv") {
+                let movies_dir = config.home_dir.join("Documents").join("movies");
+                let full_path = movies_dir.join(f);
+                let srt_path = full_path.with_extension("srt");
+                let mut nlines: Option<usize> = None;
+                if srt_path.exists() {
+                    nlines.replace(fs::read_to_string(srt_path).await?.split('\n').count());
+                }
+                let full_path: StackString = full_path.to_string_lossy().into();
+                for track in MkvTrack::get_subtitles_from_mkv(&full_path).await? {
+                    if track.track_type == Some(TrackType::Subtitles) {
+                        subtitles
+                            .entry(f.clone())
+                            .or_default()
+                            .push((track, nlines));
+                    }
+                }
+            }
+        }
+
         Ok(Self {
             local_file_list,
             file_list,
+            subtitles,
         })
     }