subtitle-merge/src/main.rs

use std::{
    ffi::{OsStr, OsString},
    fs::ReadDir,
    io,
    os::unix::{ffi::OsStrExt, process::CommandExt},
    path::{Path, PathBuf},
    process::{Command, ExitCode, ExitStatus},
};

use clap::Parser;
use itertools::Either;

#[derive(Parser)]
enum Args {
    Single(Single),
    Dir(Dir),
}

#[derive(Parser)]
struct Single {
    source: PathBuf,
    #[arg(long)]
    subtitles: Option<PathBuf>,
    #[arg(long)]
    fonts_dir: Option<PathBuf>,
    output: PathBuf,
}

#[derive(Parser)]
struct Dir {
    input_dir: PathBuf,
    output_dir: PathBuf,
}

fn find_subs(root: &Path, match_str: &OsStr) -> io::Result<Option<PathBuf>> {
    for entry in root.read_dir()? {
        let entry = entry?;
        let type_ = entry.file_type()?;
        if type_.is_dir() {
            let recursed = find_subs(&entry.path(), match_str)?;
            if recursed.is_some() {
                return Ok(recursed);
            }
        } else {
            // Assume found a file
            let path = entry.path();
            if path
                .file_stem()
                .unwrap()
                .as_bytes()
                .windows(match_str.as_bytes().len())
                .any(|subslice| subslice == match_str.as_bytes())
                && path.extension() == Some(OsStr::new("ass"))
            {
                return Ok(Some(entry.path()));
            }
        }
    }
    Ok(None)
}

fn find_files(root: &Path, extension: &OsStr) -> impl Iterator<Item = io::Result<PathBuf>> {
    struct FindFonts {
        read_dirs: Vec<ReadDir>,
        extension: OsString,
    }
    impl Iterator for FindFonts {
        type Item = io::Result<PathBuf>;

        fn next(&mut self) -> Option<Self::Item> {
            loop {
                let Some(next) = self.read_dirs.last_mut()?.next() else {
                    self.read_dirs.pop();
                    continue;
                };
                let entry = match next {
                    Err(e) => return Some(Err(e)),
                    Ok(entry) => entry,
                };
                let file_type = entry.file_type();
                let file_type = match file_type {
                    Err(e) => return Some(Err(e)),
                    Ok(file_type) => file_type,
                };
                if file_type.is_dir() {
                    match entry.path().read_dir() {
                        Err(e) => return Some(Err(e)),
                        Ok(read_dir) => {
                            self.read_dirs.push(read_dir);
                            continue;
                        }
                    }
                } else {
                    let path = entry.path();
                    if path.extension().is_some_and(|ex| ex == self.extension) {
                        return Some(Ok(path));
                    } else {
                        continue;
                    }
                }
            }
        }
    }

    match root.read_dir() {
        Ok(read_dir) => Either::Left(FindFonts {
            read_dirs: vec![read_dir],
            extension: extension.to_owned(),
        }),
        Err(e) => Either::Right(std::iter::once(Err(e))),
    }
}

fn convert_single(
    Single {
        source,
        subtitles,
        fonts_dir,
        output,
    }: Single,
) -> io::Result<Command> {
    let subtitles = if let Some(subtitles) = subtitles {
        subtitles
    } else {
        find_subs(&source.parent().unwrap(), source.file_stem().unwrap())?.unwrap()
    };
    let fonts: Vec<PathBuf> = if let Some(fonts_dir) = fonts_dir {
        fonts_dir
            .read_dir()?
            .map(|r| r.map(|entry| entry.path()))
            .collect::<io::Result<_>>()?
    } else {
        find_files(source.parent().unwrap(), OsStr::new("ttf")).collect::<io::Result<_>>()?
    };

    let mut command = Command::new("mkvmerge");
    command.arg("--output").args([
        output.into_os_string(),
        source.into_os_string(),
        subtitles.into_os_string(),
    ]);
    for font in fonts {
        command
            .args(&["--attachment-mime-type", "font/ttf", "--attach-file"])
            .arg(font.into_os_string());
    }
    Ok(command)
}

fn main() -> io::Result<ExitCode> {
    let args = Args::parse();
    match args {
        Args::Single(single) => {
            let mut command = convert_single(single)?;
            Err(command.exec())
        }
        Args::Dir(dir) => {
            let files = find_files(&dir.input_dir, OsStr::new("mkv"));
            let mut commands = Vec::new();
            for file in files {
                let input_file = file?;
                let output_file = dir
                    .output_dir
                    .join(input_file.strip_prefix(&dir.input_dir).unwrap());
                let subtitles = find_subs(&dir.input_dir, input_file.file_stem().unwrap())?;
                if subtitles.is_none() {
                    continue;
                }
                let mut command = convert_single(Single {
                    source: input_file,
                    subtitles,
                    fonts_dir: None,
                    output: output_file.clone(),
                })?;
                std::fs::create_dir_all(output_file.parent().unwrap())?;
                commands.push(command.spawn()?);
            }
            let mut status = 0;
            for mut child in commands {
                let child_status = child.wait()?;
                if status == 0 && !child_status.success() {
                    status = child_status.code().unwrap_or(1);
                }
            }
            Ok(ExitCode::from(status as u8))
        }
    }
}