api/src/video.rs

use std::{fs, io, path::Path};

use regex::Regex;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sha3::{Digest, Sha3_256};
use thiserror::Error;
use tracing::{error, info, warn};
use url::{ParseError, Url};

#[derive(Debug, Error)]
pub enum VideoError {
    #[error("Failed to parse URL: {0}")]
    UrlParse(#[from] ParseError),
    #[error("URL is an invalid YouTube URL")]
    InvalidUrl,
    #[error("Video already exists in database")]
    AlreadyExists,
    #[error("IO Error: {0}")]
    IOError(#[from] io::Error),
    #[error("Could not serialize info JSON: {0}")]
    SerializeInfoJSON(#[from] serde_json::Error),
    #[error("Failed to parse value from key '{0}'")]
    JsonKey(String),
    #[error("No video file was found for target")]
    MissingVideoFile,
}

#[derive(Debug, Default, Serialize, Deserialize, Clone)]
pub struct Video {
    pub id: i64,
    pub url: String,
    pub youtube_id: String,
    pub title: String,
    pub description: String,
    pub author: String,
    pub author_id: String,
    pub author_url: String,
    pub views: i64,
    pub upload_date: String,
    pub likes: Option<i64>,
    pub dislikes: Option<i64>,
    pub file_name: String,
    pub file_size: i64,
    pub sha256: String,
    pub thumbnail: String,
    // pub comments: Vec<Comment>
}

impl Video {
    async fn yt_dlp_task(url: &str, cookie: Option<&str>) -> Result<(), VideoError> {
        let mut args = vec![
            "--write-info-json",
            "--write-thumbnail",
            "--write-description",
            "--write-comments",
            "--no-playlist",
            "--use-postprocessor",
            "ReturnYoutubeDislike:when=pre_process",
            "-v",
            url,
        ];
        if let Some(cookie) = cookie {
            args.append(&mut vec!["--cookies", cookie]);
        }
        args.append(&mut vec![
                "-o",
                "videos/%(id)s/%(id)s.%(ext)s",
                "-f",
                "bestvideo[ext=mkv]+bestaudio[ext=m4a]/bestvideo[ext=mp4]+bestaudio[ext=m4a]/bestvideo*+bestaudio/best",
            ]);
        let mut child = tokio::process::Command::new("yt-dlp").args(args).spawn()?;
        info!("yt-dlp task invoked");
        child.wait().await?;
        info!("yt-dlp task completed successfully");

        Ok(())
    }

    fn is_url_valid(url: &Url) -> bool {
        let re = Regex::new(r"^((?:https?:)?\/\/)?((?:www|m)\.)?((?:youtube(?:-nocookie)?\.com|youtu.be))(\/(?:[\w\-]+\?v=|embed\/|live\/|v\/)?)([\w\-]+)(\S+)?$").unwrap();
        if !re.is_match(url.as_str()) {
            error!("YouTube URL RegEx match failed!");
            return false;
        }

        true
    }

    fn get_video_id(url: &Url) -> Option<String> {
        let mut pairs = url.query_pairs();
        let Some(query_v) = pairs.find(|(key, _)| key == "v") else {
            error!("Could not find 'v' query parameter in URL!");
            return None;
        };

        Some(query_v.1.to_string())
    }

    pub async fn from_url(url: &str, id: i64, cookie: Option<&str>) -> Result<Self, VideoError> {
        let url = Url::parse(url)?;
        info!("Parsed argument as URL");

        if !Self::is_url_valid(&url) {
            error!("URL is an invalid YouTube video!");
            return Err(VideoError::InvalidUrl);
        }
        let youtube_id = Self::get_video_id(&url).ok_or(VideoError::InvalidUrl)?;

        info!("URL is valid YouTube video, got ID '{youtube_id}'");

        let dir = format!("videos/{youtube_id}");
        let file_stem = format!("{dir}/{youtube_id}");
        if !Path::new(&dir).exists() {
            fs::create_dir(dir)?;
        }

        let info_json = format!("{file_stem}.info.json");
        let info_json = Path::new(&info_json);

        // ? Uploading a video doesn't mean updating it, make a PUT route for that later
        if info_json.exists() {
            warn!("Video already exists, skipping");
            return Err(VideoError::AlreadyExists);
        }

        Self::yt_dlp_task(url.as_str(), cookie).await?;

        let info: Value = serde_json::from_str(&fs::read_to_string(info_json)?)?;

        // info!("Info JSON for {youtube_id}\n{info}");

        let get_info_value = |key: &str| {
            info.get(key)
                .ok_or_else(|| VideoError::JsonKey(key.to_string()))
                .unwrap()
        };

        let get_ryd_value = |key: &str| {
            info["RYD"]["response"]
                .get(key)
                .ok_or_else(|| VideoError::JsonKey(key.to_string()))
                .unwrap()
        };

        let url = url.to_string();
        let description = format!("{file_stem}.description");
        let title = get_info_value("title").to_string();
        let author = get_info_value("uploader").to_string();
        let views = get_info_value("view_count").as_i64().unwrap_or(-1);

        // Use RYD field from info JSON
        let likes = get_ryd_value("likes").as_i64();
        let dislikes = get_ryd_value("dislikes").as_i64();

        let possible_extensions = ["mkv", "mp4", "webm"];

        let mut buffer = None;
        let mut file_name = String::new();

        for ext in possible_extensions {
            file_name = format!("{file_stem}.{ext}");
            let path = Path::new(&file_name);
            if path.exists() {
                info!("File '{}' found", path.display());
                buffer = Some(fs::read(path)?);
                break;
            }
            warn!("File '{}' not found", path.display());
        }

        let Some(buffer) = buffer else {
            error!("No video file found for {youtube_id}!");
            return Err(VideoError::MissingVideoFile);
        };

        let sha256 = format!("{:x}", Sha3_256::digest(&buffer));
        #[allow(clippy::cast_possible_wrap)]
        let file_size = buffer.len() as i64;
        let author_id = get_info_value("channel_id").to_string();
        let author_url = get_info_value("channel_url").to_string();
        let upload_date = get_info_value("upload_date").to_string();

        let thumbnail = format!("{file_stem}.webp");

        let video = Self {
            id,
            url,
            youtube_id,
            title,
            description,
            author,
            author_id,
            author_url,
            views,
            upload_date,
            likes,
            dislikes,
            file_name,
            file_size,
            sha256,
            thumbnail,
        };

        info!("Video entry so far: {video:?}");

        Ok(video)
    }
}