FlightTube/src-tauri/src/downloader.rs

//! Driving `yt-dlp` and interpreting its progress output.

/// Highest resolution available, which on YouTube means VP9 or AV1 above 1080p.
/// Audio is pinned to AAC (`m4a`) in every selector: YouTube pairs those codecs
/// with Opus, which WebKit will not decode inside an MP4 container, so taking
/// Opus would yield a silent file.
pub const FORMAT_BEST: &str = "bv*+ba[ext=m4a]/bv*+ba/b";

/// Builds a format selector for a chosen quality.
///
/// `quality` is either "best" or a maximum height in pixels ("2160", "1080", …).
/// Anything unrecognised falls back to best, so a stale stored preference can
/// never leave downloads broken.
pub fn format_selector(quality: &str) -> String {
    match quality.parse::<u32>() {
        Ok(height) if (144..=4320).contains(&height) => format!(
            "bv*[height<={h}]+ba[ext=m4a]/bv*[height<={h}]+ba/b[height<={h}]/b",
            h = height
        ),
        _ => FORMAT_BEST.to_string(),
    }
}

/// Sentinel prefix so progress lines are distinguishable from yt-dlp's ordinary
/// chatter on the same stream.
pub const PROGRESS_TEMPLATE: &str = "FTPROG %(progress.downloaded_bytes)s %(progress.total_bytes)s %(progress.speed)s %(progress.eta)s";

/// Readable, sortable filenames: upload date, then title, then the video id so
/// two videos sharing a title cannot collide.
pub const OUTPUT_TEMPLATE: &str = "%(upload_date>%Y-%m-%d)s - %(title)s [%(id)s].%(ext)s";

#[derive(Debug, Clone, PartialEq)]
pub struct Progress {
    pub downloaded: u64,
    pub total: Option<u64>,
    pub speed: Option<f64>,
    pub eta: Option<u64>,
}

impl Progress {
    pub fn pct(&self) -> Option<f64> {
        match self.total {
            Some(t) if t > 0 => Some((self.downloaded as f64 / t as f64) * 100.0),
            _ => None,
        }
    }
}

/// yt-dlp writes `NA` for values it does not know yet (notably `total_bytes`
/// before the stream is resolved), and interleaves ordinary log lines on the
/// same stream. Anything that isn't a well-formed progress line yields `None`
/// rather than an error — a garbled line must never abort a running download.
fn opt_num<T: std::str::FromStr>(tok: &str) -> Option<T> {
    if tok == "NA" || tok.is_empty() {
        None
    } else {
        tok.parse::<T>().ok()
    }
}

pub fn parse_progress_line(line: &str) -> Option<Progress> {
    let rest = line.trim().strip_prefix("FTPROG")?;
    let tokens: Vec<&str> = rest.split_whitespace().collect();
    if tokens.len() != 4 {
        return None;
    }

    // A progress line without a byte count tells us nothing; reject it.
    let downloaded: u64 = opt_num(tokens[0])?;

    Some(Progress {
        downloaded,
        total: opt_num::<u64>(tokens[1]),
        // yt-dlp emits floats like `524288.0`; parse as f64 then keep as bytes.
        speed: opt_num::<f64>(tokens[2]),
        eta: opt_num::<f64>(tokens[3]).map(|v| v as u64),
    })
}

/// Arguments for downloading one video. Kept separate from process spawning so
/// the argument construction is assertable in tests.
/// `sub_langs` is empty when subtitles are switched off, in which case none
/// are requested at all.
pub fn build_args(
    video_id: &str,
    out_template: &str,
    quality: &str,
    sub_langs: &str,
) -> Vec<String> {
    let mut args = vec![
        "-f".into(),
        format_selector(quality),
        "--merge-output-format".into(),
        "mp4".into(),
        "--no-playlist".into(),
        "--newline".into(),
        "--no-colors".into(),
        "--progress".into(),
        // Long video titles make long filenames; keep them within sane limits.
        "--trim-filenames".into(),
        "180".into(),
        "--progress-template".into(),
        PROGRESS_TEMPLATE.into(),
        // A failing subtitle must never take the video with it. Subtitles are
        // a bonus; the file is the point.
        "--ignore-errors".into(),
        "--print".into(),
        "after_move:FTPATH %(filepath)s".into(),
        "-o".into(),
        out_template.into(),
    ];

    if !sub_langs.is_empty() {
        // Subtitles come along for offline use, including YouTube's
        // auto-generated ones. WebVTT beside the video rather than muxed in:
        // WebKit reads a <track> reliably and largely ignores subtitle streams
        // inside an MP4.
        //
        // The languages are named exactly. A wildcard like "en.*" also matches
        // every machine-translated variant YouTube offers — en-en-US, en-de and
        // dozens more — and asking for all of them earns an HTTP 429.
        args.extend([
            "--write-subs".into(),
            "--write-auto-subs".into(),
            "--sub-format".into(),
            "vtt".into(),
            "--convert-subs".into(),
            "vtt".into(),
            "--sub-langs".into(),
            sub_langs.to_string(),
        ]);
    }

    args.push(format!("https://www.youtube.com/watch?v={video_id}"));
    args
}

/// Arguments for fetching only the subtitles of a video.
///
/// YouTube's HLS manifest carries audio renditions but no subtitles whatsoever,
/// so a streamed video has nothing to show unless the captions are fetched
/// separately. Auto-generated captions are included: on most videos they are
/// the only ones there are.
pub fn subs_only_args(video_id: &str, out_template: &str, sub_langs: &str) -> Vec<String> {
    vec![
        "--skip-download".into(),
        "--no-playlist".into(),
        "--no-colors".into(),
        "--ignore-errors".into(),
        "--write-subs".into(),
        "--write-auto-subs".into(),
        "--sub-format".into(),
        "vtt".into(),
        "--convert-subs".into(),
        "vtt".into(),
        "--sub-langs".into(),
        sub_langs.to_string(),
        "-o".into(),
        out_template.into(),
        format!("https://www.youtube.com/watch?v={video_id}"),
    ]
}

/// The subtitle languages to request for a preference, or empty for none.
///
/// Only the language itself and YouTube's "-orig" variant; anything broader
/// pulls in machine translations by the dozen.
pub fn sub_langs_for(pref: &str) -> String {
    if pref.is_empty() || pref == "off" {
        String::new()
    } else {
        format!("{pref},{pref}-orig")
    }
}

/// yt-dlp reports the final path via `--print after_move:`, which is more
/// reliable than guessing the extension after a merge.
pub fn parse_final_path(line: &str) -> Option<String> {
    line.trim()
        .strip_prefix("FTPATH ")
        .map(|p| p.trim().to_string())
        .filter(|p| !p.is_empty())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn subs_only_args_take_auto_generated_captions() {
        let args = subs_only_args("abc", "/tmp/%(id)s.%(ext)s", "en,en-orig");
        // Most videos have no hand-written captions at all; without this the
        // fetch comes back empty.
        assert!(args.iter().any(|a| a == "--write-auto-subs"));
        assert!(args.iter().any(|a| a == "--write-subs"));
        assert!(args.iter().any(|a| a == "--skip-download"));
        // Named exactly: a wildcard drags in dozens of machine translations.
        let langs = args.iter().position(|a| a == "--sub-langs").unwrap();
        assert_eq!(args[langs + 1], "en,en-orig");
        assert_eq!(args.last().unwrap(), "https://www.youtube.com/watch?v=abc");
    }

    #[test]
    fn subs_only_args_are_written_where_asked() {
        let args = subs_only_args("abc", "/cache/%(id)s.%(ext)s", "nl,nl-orig");
        let out = args.iter().position(|a| a == "-o").unwrap();
        assert_eq!(args[out + 1], "/cache/%(id)s.%(ext)s");
    }

    use super::*;

    #[test]
    fn parses_complete_line() {
        let p = parse_progress_line("FTPROG 1048576 10485760 524288.0 18").unwrap();
        assert_eq!(p.downloaded, 1_048_576);
        assert_eq!(p.total, Some(10_485_760));
        assert_eq!(p.speed, Some(524_288.0));
        assert_eq!(p.eta, Some(18));
    }

    #[test]
    fn handles_na_fields_before_stream_resolves() {
        let p = parse_progress_line("FTPROG 4096 NA NA NA").unwrap();
        assert_eq!(p.downloaded, 4096);
        assert_eq!(p.total, None);
        assert_eq!(p.speed, None);
        assert_eq!(p.eta, None);
    }

    #[test]
    fn tolerates_float_eta() {
        let p = parse_progress_line("FTPROG 10 100 5.5 12.0").unwrap();
        assert_eq!(p.eta, Some(12));
    }

    #[test]
    fn ignores_ordinary_yt_dlp_output() {
        assert!(parse_progress_line("[youtube] Extracting URL: https://x").is_none());
        assert!(parse_progress_line("[Merger] Merging formats into \"x.mp4\"").is_none());
        assert!(parse_progress_line("").is_none());
        assert!(parse_progress_line("[download] 50% of 10MiB").is_none());
    }

    #[test]
    fn ignores_garbage_without_panicking() {
        assert!(parse_progress_line("FTPROG").is_none());
        assert!(parse_progress_line("FTPROG a b c d").is_none());
        assert!(parse_progress_line("FTPROG 1 2").is_none());
        assert!(parse_progress_line("FTPROG 1 2 3 4 5").is_none());
        assert!(parse_progress_line("FTPROG NA NA NA NA").is_none());
    }

    #[test]
    fn computes_percent() {
        let p = parse_progress_line("FTPROG 5000 10000 1.0 1").unwrap();
        assert_eq!(p.pct(), Some(50.0));
        let q = parse_progress_line("FTPROG 5000 NA NA NA").unwrap();
        assert_eq!(q.pct(), None);
    }

    #[test]
    fn percent_guards_against_zero_total() {
        let p = Progress {
            downloaded: 5,
            total: Some(0),
            speed: None,
            eta: None,
        };
        assert_eq!(p.pct(), None);
    }

    #[test]
    fn subtitles_are_requested_including_auto_generated() {
        let args = build_args("abc", "/tmp/o.%(ext)s", "best", "en,en-orig");
        assert!(args.contains(&"--write-subs".to_string()));
        // The auto-generated track is the only one many videos have.
        assert!(args.contains(&"--write-auto-subs".to_string()));
        assert!(args.contains(&"en,en-orig".to_string()));
        // WebVTT, because that is what a <track> element can load.
        assert!(args.contains(&"vtt".to_string()));
    }

    #[test]
    fn no_preference_means_no_subtitle_requests_at_all() {
        let args = build_args("abc", "/tmp/o.%(ext)s", "best", "");
        assert!(!args.contains(&"--write-subs".to_string()));
        assert!(!args.contains(&"--write-auto-subs".to_string()));
        assert!(!args.contains(&"--sub-langs".to_string()));
    }

    #[test]
    fn a_subtitle_failure_must_not_abort_the_video() {
        // yt-dlp aborts the whole job on the first error without this.
        assert!(build_args("abc", "/tmp/o.%(ext)s", "best", "en,en-orig")
            .contains(&"--ignore-errors".to_string()));
    }

    #[test]
    fn languages_are_named_exactly_never_as_a_wildcard() {
        // "en.*" also matches en-en-US and dozens of machine translations,
        // and requesting them all earns an HTTP 429.
        let langs = sub_langs_for("nl");
        assert_eq!(langs, "nl,nl-orig");
        assert!(!langs.contains('*'));
        assert_eq!(sub_langs_for("off"), "");
        assert_eq!(sub_langs_for(""), "");
    }

    #[test]
    fn best_quality_takes_the_highest_available() {
        let args = build_args("abc123", "/tmp/out.%(ext)s", "best", "en.*");
        assert!(args.contains(&FORMAT_BEST.to_string()));
        assert!(args.contains(&"mp4".to_string()));
        assert!(args.contains(&"https://www.youtube.com/watch?v=abc123".to_string()));
        assert!(args.contains(&"--no-playlist".to_string()));
    }

    #[test]
    fn every_selector_pins_aac_audio() {
        // Opus in MP4 would be silent in WebKit, so the audio half stays m4a.
        for q in ["best", "2160", "1080", "480"] {
            assert!(format_selector(q).contains("ba[ext=m4a]"), "quality {q}");
        }
    }

    #[test]
    fn a_numeric_quality_caps_the_height() {
        let sel = format_selector("1080");
        assert!(sel.contains("height<=1080"));
        assert!(!sel.contains("height<=2160"));
        assert!(build_args("x", "o", "1080", "en.*").contains(&sel));
    }

    #[test]
    fn nonsense_or_out_of_range_quality_falls_back_to_best() {
        assert_eq!(format_selector("nonsense"), FORMAT_BEST);
        assert_eq!(format_selector(""), FORMAT_BEST);
        assert_eq!(format_selector("0"), FORMAT_BEST);
        assert_eq!(format_selector("99999"), FORMAT_BEST);
        assert_eq!(format_selector("best"), FORMAT_BEST);
    }

    #[test]
    fn output_template_is_date_then_title_then_id() {
        assert!(OUTPUT_TEMPLATE.starts_with("%(upload_date>%Y-%m-%d)s - %(title)s"));
        assert!(OUTPUT_TEMPLATE.contains("[%(id)s]"));
    }

    #[test]
    fn extracts_final_path() {
        assert_eq!(
            parse_final_path("FTPATH /Users/x/Movies/FlightTube/abc.mp4").unwrap(),
            "/Users/x/Movies/FlightTube/abc.mp4"
        );
        assert!(parse_final_path("FTPATH ").is_none());
        assert!(parse_final_path("[download] done").is_none());
    }
}