zero4rs 2.0.0

zero4rs is a powerful, pragmatic, and extremely fast web framework for Rust
Documentation
pub mod args;
pub mod stream_filter;
pub mod stream_serializer;
pub mod video_serializer;

use std::path::PathBuf;

use anyhow::{Context, Result};

use rustube::video_info::player_response::streaming_data::Quality;
use rustube::Callback;
use rustube::{Error, Id, IdBuf, Stream, Video, VideoFetcher, VideoInfo};

use crate::core::youtube::args::command;
use crate::core::youtube::args::download;
use crate::core::youtube::args::output::OutputArgs;
use crate::core::youtube::args::output_format::OutputFormat;
use crate::core::youtube::args::output_level::OutputLevel;
use crate::core::youtube::stream_filter::StreamFilter;
use crate::core::youtube::video_serializer::VideoSerializer;

pub async fn check(args: command::CheckArgs) -> Result<()> {
    let output = check_by_id(&args.identifier.id_string()).await?;
    log::info!("{output}");
    Ok(())
}

pub async fn check_by_id(id: &str) -> Result<String> {
    let id = Id::from_raw(id)?.into_owned();

    let stream_filter = StreamFilter {
        best_quality: true,
        worst_quality: false,
        no_video: false,
        no_audio: false,
        ignore_missing_video: false,
        ignore_missing_audio: false,
        quality: None,
        video_quality: None,
        audio_quality: None,
    };

    let output = OutputArgs {
        output_format: OutputFormat::Yaml,
        output_level: OutputLevel::URL
            | OutputLevel::GENERAL
            | OutputLevel::VIDEO_TRACK
            | OutputLevel::AUDIO_TRACK,
    };

    let (video_info, streams) = get_streams(id, &stream_filter).await?;

    let video_serializer = VideoSerializer::new(video_info, streams, output.output_level);

    let output = output.output_format.serialize_output(&video_serializer)?;

    Ok(output)
}

pub async fn content_length(id: &str) -> Result<(String, u64, Quality)> {
    let id = Id::from_raw(id)?.into_owned();

    let stream_filter = StreamFilter {
        best_quality: true,
        worst_quality: false,
        no_video: false,
        no_audio: false,
        ignore_missing_video: false,
        ignore_missing_audio: false,
        quality: None,
        video_quality: None,
        audio_quality: None,
    };

    let (_, stream) = get_stream(id.as_owned(), stream_filter).await?;
    let content_length = stream.content_length().await?;

    log::info!(
        "youtube-content_length: id={}, content_length={}",
        id,
        content_length
    );

    Ok((
        stream.video_details.title.clone(),
        content_length,
        stream.quality,
    ))
}

pub async fn video_info(id: &str) -> Result<VideoInfo> {
    let id = Id::from_raw(id)?.into_owned();

    let stream_filter = StreamFilter {
        best_quality: true,
        worst_quality: false,
        no_video: false,
        no_audio: false,
        ignore_missing_video: false,
        ignore_missing_audio: false,
        quality: None,
        video_quality: None,
        audio_quality: None,
    };

    let (video_info, _) = get_stream(id.as_owned(), stream_filter).await?;

    Ok(video_info)
}

pub async fn best_quality(id: &str) -> Result<Stream> {
    let id = Id::from_raw(id)?.into_owned();

    let stream_filter = StreamFilter {
        best_quality: true,
        worst_quality: false,
        no_video: false,
        no_audio: false,
        ignore_missing_video: false,
        ignore_missing_audio: false,
        quality: None,
        video_quality: None,
        audio_quality: None,
    };

    let (_, stream) = get_stream(id.as_owned(), stream_filter).await?;

    Ok(stream)
}

pub async fn download_by_id(id: &str, dir: &str) -> Result<PathBuf> {
    let id = Id::from_raw(id)?.into_owned();

    let stream_filter = StreamFilter {
        best_quality: true,
        worst_quality: false,
        no_video: false,
        no_audio: false,
        ignore_missing_video: false,
        ignore_missing_audio: false,
        quality: None,
        video_quality: None,
        audio_quality: None,
    };

    let (_video_info, stream) = get_stream(id.as_owned(), stream_filter).await?;
    let content_length = stream.content_length().await?;
    let ext_name = stream.mime.subtype().as_str();

    let download_path = download_path(None, ext_name, Some(dir.parse()?), None, id);

    log::info!(
        "youtube-download_by_id: content_length={}, download_path={:?}",
        content_length,
        download_path.as_os_str()
    );

    stream.download_to(&download_path).await?;

    Ok(download_path)
}

pub async fn download_async(
    id_str: &str,
    dir: &str,
    prefix: Option<&String>,
) -> Result<(String, u64, Quality)> {
    let id = Id::from_raw(id_str)?.into_owned();

    let stream_filter = StreamFilter {
        best_quality: true,
        worst_quality: false,
        no_video: false,
        no_audio: false,
        ignore_missing_video: false,
        ignore_missing_audio: false,
        quality: None,
        video_quality: None,
        audio_quality: None,
    };

    let (_video_info, stream) = get_stream(id.as_owned(), stream_filter).await?;
    let content_length = stream.content_length().await?;
    let ext_name = stream.mime.subtype().as_str();
    let title = stream.video_details.title.clone();
    let quality = stream.quality;
    let title_escape = eacape_title(&title, prefix);

    let download_path = download_path(None, ext_name, Some(dir.parse()?), Some(&title_escape), id);

    tokio::spawn(async move {
        match stream.download_to(&download_path).await {
            Ok(()) => {
                log::info!(
                    "youtube-download_async-completed: content_length={}, title_escape={}, download_path={:?}",
                    crate::commons::format_bytes_size(content_length as usize),
                    title_escape,
                    &download_path
                );
            }
            Err(e) => {
                log::error!(
                    "youtube-download_async: path={:?}, error={:?}",
                    &download_path,
                    e
                );
            }
        }
    });

    log::info!(
        "youtube-download_async: id={}, content_length={}, title={}, dir={}",
        id_str,
        content_length,
        title,
        dir
    );

    Ok((title, content_length, quality))
}

pub async fn download(args: download::DownloadArgs) -> Result<()> {
    let id = args.identifier.id()?;
    let (video_info, stream) = get_stream(id.as_owned(), args.stream_filter).await?;
    let download_path = download_path(
        args.filename,
        stream.mime.subtype().as_str(),
        args.dir,
        None,
        id,
    );
    let title = format!("{}.mp4", args.identifier.id_string());
    let content_length = stream.content_length().await?;

    let pb = indicatif::ProgressBar::new(content_length);

    pb.set_style(indicatif::ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
        .unwrap()
        .with_key("eta", move |state: &indicatif::ProgressState, w: &mut dyn std::fmt::Write| write!(w, "{:.1}s {}", state.eta().as_secs_f64(), title).unwrap())
        .progress_chars("#>-"));

    let callback = Callback::new().connect_on_progress_closure(move |cargs| {
        if cargs.current_chunk as u64 <= content_length {
            pb.set_position(cargs.current_chunk as u64);
        } else {
            pb.finish_with_message("downloaded");
        }
    });

    stream
        .download_to_with_callback(&download_path, callback)
        .await?;

    // download completed
    let video_serializer = VideoSerializer::new(
        video_info,
        std::iter::once(stream),
        args.output.output_level,
    );

    let output = args
        .output
        .output_format
        .serialize_output(&video_serializer)
        .unwrap();

    log::info!("{output}");

    Ok(())
}

pub async fn get_stream(
    id: IdBuf,
    stream_filter: stream_filter::StreamFilter,
) -> Result<(VideoInfo, Stream)> {
    let (video_info, streams) = get_streams(id, &stream_filter).await?;

    let stream = streams
        .max_by(|lhs, rhs| stream_filter.max_stream(lhs, rhs))
        .ok_or(Error::NoStreams)
        .context("There are no streams, that match all your criteria")?;

    Ok((video_info, stream))
}

pub async fn get_streams(
    id: IdBuf,
    stream_filter: &'_ stream_filter::StreamFilter,
) -> Result<(VideoInfo, impl Iterator<Item = Stream> + '_)> {
    let (video_info, streams) = get_video(id).await?.into_parts();

    let streams = streams
        .into_iter()
        .filter(move |stream| stream_filter.stream_matches(stream));

    Ok((video_info, streams))
}

pub async fn get_video(id: IdBuf) -> Result<Video> {
    VideoFetcher::from_id(id)?
        .fetch()
        .await
        .context("Could not fetch the video information")?
        .descramble()
        .context("Could not descramble the video information")
}

pub fn download_path(
    filename: Option<PathBuf>,
    extension: &str,
    dir: Option<PathBuf>,
    _title: Option<&str>,
    video_id: Id<'_>,
) -> PathBuf {
    let video_id = if let Some(title) = _title {
        format!("{}-{}", title, video_id.as_str())
    } else {
        video_id.as_str().to_string()
    };

    let filename = filename.unwrap_or_else(|| format!("{}.{}", video_id, extension).into());

    let mut path = dir.unwrap_or_default();

    path.push(filename);
    path
}

pub fn eacape_title(title: &str, prefix: Option<&String>) -> String {
    // 匹配除了 a-zA-Z0-9 和中文之外的所有字符
    let re = regex::Regex::new(r#"[^-a-zA-Z0-9\u4e00-\u9fa5]+"#).unwrap();
    let title = re.replace_all(title, " ");

    // 匹配2个或多个空格
    let re = regex::Regex::new(r#"\s{2,}"#).unwrap();
    let title = re.replace_all(&title, " ");

    // 匹配 -
    let re = regex::Regex::new(r#" - "#).unwrap();
    let title = re.replace_all(&title, "-");

    // 匹配任意空格
    let re = regex::Regex::new(r#"\s+"#).unwrap();
    let title = re.replace_all(title.trim(), "_");

    if let Some(prefix) = prefix {
        format!("{}{}", prefix, title)
    } else {
        title.to_string()
    }
}