crabidy/ytdy/src/extract.rs

//! The extraction seam between the provider and YouTube
//! (see `architecture/youtube-rustypipe.md` D1).
//!
//! Everything the provider needs from YouTube goes through the
//! [`Extract`] trait: the real implementation wraps a pure-Rust
//! [`rustypipe`] Innertube client (no subprocess, no Python), tests use
//! a fake. All errors are typed; messages may name videos, playlists,
//! and public API endpoints — never cookie values or auth headers.

use std::fmt::Debug;
use std::path::{Path, PathBuf};
use std::time::Duration;

use async_trait::async_trait;
use rustypipe::client::RustyPipe;
use rustypipe::model::VideoItem;
use tracing::{debug, warn};

/// Errors from the extractor.
#[derive(Debug, thiserror::Error)]
pub enum ExtractError {
    /// The client could not be built (storage directory, TLS backend).
    #[error("cannot build the youtube client: {0}")]
    Client(String),
    /// A YouTube request failed (network, extraction, API change).
    #[error("youtube request failed: {0}")]
    Fetch(String),
    /// The video exists but offers no audio streams.
    #[error("video has no audio streams")]
    NoAudio,
}

/// One video as the provider sees it.
#[derive(Clone, Debug)]
pub struct VideoEntry {
    /// YouTube video id (the track path's last segment).
    pub id: String,
    pub title: String,
    /// Channel name; the wire track's artist.
    pub artist: Option<String>,
    /// Duration in seconds; `None` for livestreams.
    pub duration: Option<u32>,
}

/// One playlist of the logged-in user.
#[derive(Clone, Debug)]
pub struct PlaylistEntry {
    /// YouTube playlist id (the node path's last segment).
    pub id: String,
    pub title: String,
}

/// What the provider needs from YouTube. Implementations must be
/// side-effect free towards the provider: every call is independent,
/// there is no session state beyond the login performed at init.
#[async_trait]
pub trait Extract: Send + Sync + Debug {
    /// The top `limit` video results for a search query, result order.
    async fn search_videos(
        &self,
        query: &str,
        limit: usize,
    ) -> Result<Vec<VideoEntry>, ExtractError>;
    /// Metadata of one video.
    async fn video(&self, id: &str) -> Result<VideoEntry, ExtractError>;
    /// The playable audio stream URL of one video — `audio/mp4` (AAC)
    /// preferred, since the local player cannot decode Opus
    /// (architecture/youtube-rustypipe.md D2).
    async fn audio_stream_url(&self, id: &str) -> Result<String, ExtractError>;
    /// The logged-in user's saved playlists (first page).
    async fn saved_playlists(&self) -> Result<Vec<PlaylistEntry>, ExtractError>;
    /// One playlist's name and up to `limit` videos (bounded
    /// pagination).
    async fn playlist_videos(
        &self,
        id: &str,
        limit: usize,
    ) -> Result<(String, Vec<VideoEntry>), ExtractError>;
}

/// The real extractor: one [`RustyPipe`] Innertube client.
pub struct RustyPipeExtractor {
    rp: RustyPipe,
}

impl Debug for RustyPipeExtractor {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("RustyPipeExtractor").finish_non_exhaustive()
    }
}

/// Maps a rustypipe failure to [`ExtractError::Fetch`]. rustypipe's
/// errors are typed (HTTP status, extraction, auth) and never carry
/// cookie values.
fn fetch_err(err: rustypipe::error::Error) -> ExtractError {
    ExtractError::Fetch(err.to_string())
}

fn video_entry(item: VideoItem) -> VideoEntry {
    VideoEntry {
        id: item.id,
        title: item.name,
        artist: item.channel.map(|channel| channel.name),
        duration: item.duration,
    }
}

/// Picks the audio stream to play from `(mime, average_bitrate)` pairs:
/// the highest-bitrate `audio/mp4` (AAC — the local rodio/symphonia
/// player decodes it), falling back to the highest-bitrate stream of
/// any type when no mp4 exists. Returns the index into the input.
fn pick_audio_stream(streams: &[(String, u32)]) -> Option<usize> {
    let best = |mp4_only: bool| {
        streams
            .iter()
            .enumerate()
            .filter(|(_, (mime, _))| !mp4_only || mime.starts_with("audio/mp4"))
            .max_by_key(|(_, (_, bitrate))| *bitrate)
            .map(|(index, _)| index)
    };
    best(true).or_else(|| best(false))
}

impl RustyPipeExtractor {
    /// Builds the client. `storage_dir` holds rustypipe's cache file
    /// (client state and the rotated auth cookie — as secret as the
    /// cookies file itself); `timeout` bounds every YouTube request.
    /// `botguard_bin` optionally names a `rustypipe-botguard` binary for
    /// PO-token attestation (rustypipe also auto-detects it on PATH).
    pub fn new(
        storage_dir: PathBuf,
        timeout: Duration,
        botguard_bin: Option<&Path>,
    ) -> Result<Self, ExtractError> {
        std::fs::create_dir_all(&storage_dir)
            .map_err(|err| ExtractError::Client(format!("storage directory: {err}")))?;
        let mut builder = RustyPipe::builder()
            .storage_dir(storage_dir)
            .timeout(timeout);
        if let Some(bin) = botguard_bin {
            builder = builder.botguard_bin(bin.as_os_str().to_owned());
        }
        let rp = builder
            .build()
            .map_err(|err| ExtractError::Client(err.to_string()))?;
        Ok(Self { rp })
    }

    /// Logs in from a Netscape `cookies.txt` export, cache first
    /// (architecture/youtube-rustypipe.md D3): a still-valid rotated
    /// cookie in the rustypipe cache wins over re-reading the
    /// stale-prone file. Returns whether the client is authenticated.
    /// Never fails — anything short of a working login degrades to
    /// logged-out with a warning naming only the file *path*.
    pub async fn login(&self, cookies_file: Option<&Path>) -> bool {
        let Some(path) = cookies_file else {
            return false;
        };
        // The cached cookie is refreshed by rustypipe on use and
        // outlives the original export (YouTube rotates cookies).
        if self.rp.user_auth_check_cookie().await.is_ok() {
            debug!("youtube login restored from the rustypipe cache");
            return true;
        }
        let contents = match std::fs::read_to_string(path) {
            Ok(contents) => contents,
            Err(err) => {
                warn!(
                    cookies = %path.display(),
                    "cookies file not readable, running logged out: {err}"
                );
                return false;
            }
        };
        match self.rp.user_auth_set_cookie_txt(&contents).await {
            Ok(()) => {
                debug!(cookies = %path.display(), "youtube login succeeded");
                true
            }
            Err(err) => {
                warn!(
                    cookies = %path.display(),
                    "youtube login failed, running logged out: {err}"
                );
                false
            }
        }
    }
}

#[async_trait]
impl Extract for RustyPipeExtractor {
    async fn search_videos(
        &self,
        query: &str,
        limit: usize,
    ) -> Result<Vec<VideoEntry>, ExtractError> {
        let result = self
            .rp
            .query()
            .search::<VideoItem, _>(query)
            .await
            .map_err(fetch_err)?;
        Ok(result
            .items
            .items
            .into_iter()
            .take(limit)
            .map(video_entry)
            .collect())
    }

    async fn video(&self, id: &str) -> Result<VideoEntry, ExtractError> {
        // `player` (not `video_details`) because it carries the
        // duration — and it is the same call the stream fetch uses.
        let player = self.rp.query().player(id).await.map_err(fetch_err)?;
        let details = player.details;
        Ok(VideoEntry {
            id: details.id,
            title: details.name.unwrap_or_default(),
            artist: details.channel_name,
            duration: (details.duration > 0).then_some(details.duration),
        })
    }

    async fn audio_stream_url(&self, id: &str) -> Result<String, ExtractError> {
        let player = self.rp.query().player(id).await.map_err(fetch_err)?;
        let candidates: Vec<(String, u32)> = player
            .audio_streams
            .iter()
            .map(|stream| (stream.mime.clone(), stream.average_bitrate))
            .collect();
        let Some(index) = pick_audio_stream(&candidates) else {
            return Err(ExtractError::NoAudio);
        };
        let stream = &player.audio_streams[index];
        if !stream.mime.starts_with("audio/mp4") {
            // Playable URL, but the local player has no Opus decoder —
            // hand it out anyway (a future player may cope) and say so.
            warn!(video = id, mime = %stream.mime, "no mp4 audio stream; local decoding may fail");
        }
        Ok(stream.url.clone())
    }

    async fn saved_playlists(&self) -> Result<Vec<PlaylistEntry>, ExtractError> {
        let playlists = self.rp.query().saved_playlists().await.map_err(fetch_err)?;
        Ok(playlists
            .items
            .into_iter()
            .map(|item| PlaylistEntry {
                id: item.id,
                title: item.name,
            })
            .collect())
    }

    async fn playlist_videos(
        &self,
        id: &str,
        limit: usize,
    ) -> Result<(String, Vec<VideoEntry>), ExtractError> {
        let mut playlist = self.rp.query().playlist(id).await.map_err(fetch_err)?;
        playlist
            .videos
            .extend_limit(self.rp.query(), limit)
            .await
            .map_err(fetch_err)?;
        let videos = playlist
            .videos
            .items
            .into_iter()
            .take(limit)
            .map(video_entry)
            .collect();
        Ok((playlist.name, videos))
    }
}

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

    #[test]
    fn stream_pick_prefers_mp4_then_bitrate() {
        let streams = |list: &[(&str, u32)]| -> Vec<(String, u32)> {
            list.iter().map(|(m, b)| (m.to_string(), *b)).collect()
        };
        // The best mp4 wins even against a higher-bitrate opus.
        let mixed = streams(&[
            ("audio/webm; codecs=\"opus\"", 128_000),
            ("audio/mp4; codecs=\"mp4a.40.2\"", 96_000),
            ("audio/mp4; codecs=\"mp4a.40.5\"", 48_000),
        ]);
        assert_eq!(pick_audio_stream(&mixed), Some(1));
        // Without mp4, the highest bitrate of anything is used.
        let opus_only = streams(&[
            ("audio/webm; codecs=\"opus\"", 64_000),
            ("audio/webm; codecs=\"opus\"", 128_000),
        ]);
        assert_eq!(pick_audio_stream(&opus_only), Some(1));
        assert_eq!(pick_audio_stream(&[]), None);
    }

    #[tokio::test]
    async fn login_without_a_readable_file_degrades_quietly() {
        // Network-free: no cached cookie (fresh storage dir) and a
        // missing file short-circuit before any request.
        let dir = tempfile::TempDir::new().expect("tempdir");
        let extractor =
            RustyPipeExtractor::new(dir.path().join("rustypipe"), Duration::from_secs(5), None)
                .expect("client builds");
        assert!(!extractor.login(None).await);
        assert!(!extractor.login(Some(&dir.path().join("gone.txt"))).await);
    }
}