psyche_subtitle_toolkit/media/
mkv.rs

1use std::path::{Path, PathBuf};
2use std::process::Stdio;
3
4use serde::Deserialize;
5use tokio::process::Command;
6
7use crate::error::{Result, SubtitleToolkitError};
8
9/// Parsed MKV container information from `mkvmerge -J`.
10#[derive(Debug, Clone, Deserialize)]
11pub struct MkvInfo {
12    /// List of tracks in the MKV.
13    pub tracks: Vec<MkvTrack>,
14}
15
16/// A single track inside an MKV container.
17#[derive(Debug, Clone, Deserialize)]
18pub struct MkvTrack {
19    /// Track ID (0-based).
20    pub id: u64,
21    /// Track type: `"video"`, `"audio"`, or `"subtitles"`.
22    #[serde(rename = "type")]
23    pub track_type: String,
24    /// Codec name (e.g. `"SubStationAlpha"`, `"AVC/H.264/MPEG-4p10"`).
25    pub codec: Option<String>,
26    /// Track metadata properties.
27    pub properties: MkvTrackProperties,
28}
29
30/// Metadata properties for an MKV track.
31#[derive(Debug, Clone, Deserialize)]
32pub struct MkvTrackProperties {
33    /// BCP 47 language code (e.g. `"eng"`, `"jpn"`).
34    pub language: Option<String>,
35    /// User-defined track name.
36    pub track_name: Option<String>,
37}
38
39/// Subtitle format detected in an MKV track.
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
41pub enum SubtitleFormat {
42    /// ASS/SSA subtitle format.
43    Ass,
44    /// SRT (SubRip) subtitle format.
45    Srt,
46    /// WebVTT subtitle format.
47    Vtt,
48}
49
50impl MkvTrack {
51    /// Returns `true` if this track is an ASS/SSA subtitle.
52    pub fn is_ass_subtitle(&self) -> bool {
53        self.track_type == "subtitles"
54            && self.codec.as_deref().is_some_and(|codec| {
55                codec.eq_ignore_ascii_case("SubStationAlpha")
56                    || codec.eq_ignore_ascii_case("AdvancedSubStationAlpha")
57            })
58    }
59
60    /// Returns `true` if this track is an SRT (SubRip) subtitle.
61    pub fn is_srt_subtitle(&self) -> bool {
62        self.track_type == "subtitles"
63            && self.codec.as_deref().is_some_and(|codec| {
64                codec.eq_ignore_ascii_case("SubRip")
65                    || codec.eq_ignore_ascii_case("SRT")
66                    || codec.eq_ignore_ascii_case("SubRip/SRT")
67            })
68    }
69
70    /// Returns `true` if this track is a WebVTT subtitle.
71    pub fn is_vtt_subtitle(&self) -> bool {
72        self.track_type == "subtitles"
73            && self.codec.as_deref().is_some_and(|codec| {
74                codec.eq_ignore_ascii_case("WebVTT")
75                    || codec.eq_ignore_ascii_case("VTT")
76            })
77    }
78
79    /// Returns the subtitle format if this track is a supported subtitle.
80    pub fn subtitle_format(&self) -> Option<SubtitleFormat> {
81        if self.is_ass_subtitle() {
82            Some(SubtitleFormat::Ass)
83        } else if self.is_srt_subtitle() {
84            Some(SubtitleFormat::Srt)
85        } else if self.is_vtt_subtitle() {
86            Some(SubtitleFormat::Vtt)
87        } else {
88            None
89        }
90    }
91}
92
93/// Discover MKV files at the given path.
94///
95/// If `input` is a file, returns it as a single-element list.
96/// If `input` is a directory, scans for `.mkv` files (non-recursive).
97pub async fn discover_mkv_files(input: &Path) -> Result<Vec<PathBuf>> {
98    if input.is_file() {
99        if is_mkv(input) {
100            return Ok(vec![input.to_path_buf()]);
101        }
102        return Err(SubtitleToolkitError::NoMkvFiles {
103            path: input.to_path_buf(),
104        });
105    }
106
107    let mut files = Vec::new();
108    let mut entries = tokio::fs::read_dir(input).await?;
109    while let Some(entry) = entries.next_entry().await? {
110        let path = entry.path();
111        if path.is_file() && is_mkv(&path) {
112            files.push(path);
113        }
114    }
115    files.sort();
116
117    if files.is_empty() {
118        return Err(SubtitleToolkitError::NoMkvFiles {
119            path: input.to_path_buf(),
120        });
121    }
122
123    Ok(files)
124}
125
126/// Inspect an MKV file and return track information.
127///
128/// Requires `mkvmerge` to be installed and in PATH.
129pub async fn inspect_mkv(path: &Path) -> Result<MkvInfo> {
130    let output = run_output("mkvmerge", ["-J".into(), path.as_os_str().into()]).await?;
131    Ok(serde_json::from_slice(&output)?)
132}
133
134/// Select the first ASS subtitle track, or a specific track by ID.
135///
136/// Returns `None` if no matching ASS track is found.
137pub fn select_ass_track(info: &MkvInfo, requested_track: Option<u64>) -> Option<&MkvTrack> {
138    if let Some(track_id) = requested_track {
139        return info
140            .tracks
141            .iter()
142            .find(|track| track.id == track_id && track.is_ass_subtitle());
143    }
144
145    info.tracks.iter().find(|track| track.is_ass_subtitle())
146}
147
148/// Select the first supported subtitle track (ASS, SRT, or VTT), or a specific track by ID.
149///
150/// Returns the track and its detected format. Prefers ASS > SRT > VTT when multiple exist.
151/// Returns `None` if no matching subtitle track is found.
152pub fn select_subtitle_track(
153    info: &MkvInfo,
154    requested_track: Option<u64>,
155) -> Option<(&MkvTrack, SubtitleFormat)> {
156    if let Some(track_id) = requested_track {
157        let track = info.tracks.iter().find(|t| t.id == track_id)?;
158        let format = track.subtitle_format()?;
159        return Some((track, format));
160    }
161
162    // Prefer ASS > SRT > VTT
163    if let Some(track) = info.tracks.iter().find(|t| t.is_ass_subtitle()) {
164        return Some((track, SubtitleFormat::Ass));
165    }
166
167    if let Some(track) = info.tracks.iter().find(|t| t.is_srt_subtitle()) {
168        return Some((track, SubtitleFormat::Srt));
169    }
170
171    info.tracks
172        .iter()
173        .find(|t| t.is_vtt_subtitle())
174        .map(|track| (track, SubtitleFormat::Vtt))
175}
176
177/// Extract a subtitle track from an MKV file.
178///
179/// Requires `mkvextract` to be installed and in PATH.
180pub async fn extract_subtitle(input: &Path, track_id: u64, output: &Path) -> Result<()> {
181    let selector = format!("{track_id}:{}", output.display());
182    run_status(
183        "mkvextract",
184        [input.as_os_str().into(), "tracks".into(), selector.into()],
185    )
186    .await
187}
188
189/// Mux a subtitle file into an MKV, replacing the specified track.
190///
191/// Writes to a temporary file first, then renames over the original for safety.
192/// Requires `mkvmerge` to be installed and in PATH.
193pub async fn mux_subtitle_in_place(
194    input: &Path,
195    replaced_track_id: u64,
196    subtitle: &Path,
197    language: &str,
198) -> Result<()> {
199    let parent = input.parent().unwrap_or_else(|| Path::new("."));
200    let stem = input
201        .file_stem()
202        .and_then(|stem| stem.to_str())
203        .unwrap_or("output");
204    let temp_output = parent.join(format!(".{stem}.psyche-subtitle-toolkit.tmp.mkv"));
205    let subtitle_arg = subtitle.as_os_str();
206    let language_arg = format!("0:{language}");
207    let subtitle_tracks_arg = format!("!{replaced_track_id}");
208
209    run_status(
210        "mkvmerge",
211        [
212            "-o".into(),
213            temp_output.as_os_str().into(),
214            "--subtitle-tracks".into(),
215            subtitle_tracks_arg.into(),
216            input.as_os_str().into(),
217            "--language".into(),
218            language_arg.into(),
219            subtitle_arg.into(),
220        ],
221    )
222    .await?;
223
224    tokio::fs::rename(temp_output, input).await?;
225    Ok(())
226}
227
228fn is_mkv(path: &Path) -> bool {
229    path.extension()
230        .and_then(|extension| extension.to_str())
231        .is_some_and(|extension| extension.eq_ignore_ascii_case("mkv"))
232}
233
234async fn run_output(
235    program: &'static str,
236    args: impl IntoIterator<Item = std::ffi::OsString>,
237) -> Result<Vec<u8>> {
238    let output = Command::new(program)
239        .args(args)
240        .stdout(Stdio::piped())
241        .stderr(Stdio::piped())
242        .output()
243        .await
244        .map_err(|error| {
245            if error.kind() == std::io::ErrorKind::NotFound {
246                SubtitleToolkitError::MissingTool { tool: program }
247            } else {
248                SubtitleToolkitError::Io(error)
249            }
250        })?;
251
252    if !output.status.success() {
253        return Err(SubtitleToolkitError::CommandFailed {
254            program,
255            status: output.status.to_string(),
256            stderr: String::from_utf8_lossy(&output.stderr).into_owned(),
257        });
258    }
259
260    Ok(output.stdout)
261}
262
263async fn run_status(
264    program: &'static str,
265    args: impl IntoIterator<Item = std::ffi::OsString>,
266) -> Result<()> {
267    run_output(program, args).await.map(|_| ())
268}
psyche_subtitle_toolkit/media/mkv.rs

psyche_subtitle_toolkit/media/
mkv.rs
