TelegramTwitterMediaBot/crates/x-media/src/site/pixiv/api.rs

//! Native pixiv app-API client (replaces pixiv3-rs).
//!
//! Token exchange against `oauth.secure.pixiv.net` and illust detail against
//! `app-api.pixiv.net`, deserialized with the kept `model.rs` types.

use super::interface::Illustration;
use super::model::{IllustrationModel, TypeModel, UgoiraMetadataModel};
use crate::media::Media;
use crate::site::FetchError;
use std::env;
use std::fmt;
use std::io::Read;
use std::sync::LazyLock;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, SystemTime};

const AUTH_TOKEN_URL: &str = "https://oauth.secure.pixiv.net/auth/token";
const APP_API_URL: &str = "https://app-api.pixiv.net";
const CLIENT_ID: &str = "MOBrBDS8blbauoSck0ZfDbtuzpyT";
const CLIENT_SECRET: &str = "lsACyCD94FhDUtGTXi3QzcFE2uU1hqtDaKeqrdwj";
const AUTH_USER_AGENT: &str = "PixivAndroidApp/5.0.234 (Android 11; Pixel 5)";
const APP_USER_AGENT: &str = "PixivIOSApp/7.13.3 (iOS 14.6; iPhone13,2)";
/// Token refresh safe margin (seconds).
const TOKEN_REFRESH_SAFE_MARGIN: u64 = 300;

#[derive(Debug)]
pub enum PixivError {
    /// No refresh token available (PIXIV_REFRESH_TOKEN unset).
    NoAuth,
    Http(reqwest::Error),
    Json(serde_json::Error),
    Api(String),
}

impl fmt::Display for PixivError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            PixivError::NoAuth => write!(f, "pixiv: no authentication"),
            PixivError::Http(e) => write!(f, "pixiv http error: {e}"),
            PixivError::Json(e) => write!(f, "pixiv json error: {e}"),
            PixivError::Api(message) => write!(f, "pixiv api error: {message}"),
        }
    }
}

impl std::error::Error for PixivError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            PixivError::Http(e) => Some(e),
            PixivError::Json(e) => Some(e),
            _ => None,
        }
    }
}

impl From<reqwest::Error> for PixivError {
    fn from(e: reqwest::Error) -> Self {
        PixivError::Http(e)
    }
}

impl From<serde_json::Error> for PixivError {
    fn from(e: serde_json::Error) -> Self {
        PixivError::Json(e)
    }
}

/// Native pixiv app-API client.
pub struct PixivAPI {
    refresh_token: String,
    access_token: tokio::sync::Mutex<Option<(String, SystemTime)>>,
}

impl PixivAPI {
    pub fn new(refresh_token: String) -> Self {
        Self {
            refresh_token,
            access_token: tokio::sync::Mutex::new(None),
        }
    }

    /// Returns a valid access token, exchanging the refresh token when none
    /// is cached or it has expired.
    pub async fn get_access_token(&self) -> Result<String, PixivError> {
        let mut guard = self.access_token.lock().await;
        if let Some((token, expires_at)) = guard.as_ref()
            && *expires_at > SystemTime::now()
        {
            return Ok(token.clone());
        }
        let response = crate::site::CLIENT
            .post(AUTH_TOKEN_URL)
            .form(&[
                ("client_id", CLIENT_ID),
                ("client_secret", CLIENT_SECRET),
                ("grant_type", "refresh_token"),
                ("include_policy", "true"),
                ("refresh_token", &self.refresh_token),
            ])
            .header("User-Agent", AUTH_USER_AGENT)
            .send()
            .await?;
        let json: serde_json::Value = serde_json::from_str(&response.text().await?)?;
        let access_token = json
            .get("access_token")
            .and_then(|v| v.as_str())
            .ok_or_else(|| {
                let message = json
                    .get("error")
                    .and_then(|v| v.as_str())
                    .unwrap_or("invalid token response");
                PixivError::Api(message.to_string())
            })?
            .to_string();
        let expires_in = json
            .get("expires_in")
            .and_then(|v| v.as_u64())
            .filter(|&sec| sec > 0)
            .unwrap_or(3600);
        let expires_at = SystemTime::now()
            + Duration::from_secs(expires_in.saturating_sub(TOKEN_REFRESH_SAFE_MARGIN));
        *guard = Some((access_token.clone(), expires_at));
        Ok(access_token)
    }

    /// Fetches illust detail from the app API.
    pub async fn illust_detail(&self, illust_id: u64) -> Result<IllustrationModel, PixivError> {
        let access_token = self.get_access_token().await?;
        let response = crate::site::CLIENT
            .get(format!(
                "{APP_API_URL}/v1/illust/detail?illust_id={illust_id}"
            ))
            .header("app-os", "ios")
            .header("app-os-version", "14.6")
            .header("User-Agent", APP_USER_AGENT)
            .bearer_auth(access_token)
            .send()
            .await?;
        if !response.status().is_success() {
            return Err(PixivError::Api(format!("status {}", response.status())));
        }
        let json: serde_json::Value = serde_json::from_str(&response.text().await?)?;
        if json.get("error").is_some() {
            let message = json
                .get("message")
                .and_then(|v| v.as_str())
                .unwrap_or("illust detail failed");
            return Err(PixivError::Api(message.to_string()));
        }
        let illust = json
            .get("illust")
            .ok_or_else(|| PixivError::Api("missing illust in response".to_string()))?;
        Ok(serde_json::from_value(illust.clone())?)
    }

    pub async fn fetch(&self, illust_id: u64) -> Result<Illustration, FetchError> {
        let model = self.illust_detail(illust_id).await?;
        let mut illustration = Illustration::from_model(&model);
        if matches!(&model.r#type, TypeModel::Ugoira) {
            // Real ugoira support: download the frame zip and encode an MP4.
            // Without ffmpeg (or on encode failure) the post stays
            // unsupported (empty media, like Python).
            match self.ugoira_video(illust_id).await {
                Ok(Some((mp4_path, _keep_alive))) => {
                    illustration.media.push(Media::Video {
                        title: None,
                        url: mp4_path,
                        thumbnail_url: model.image_urls.medium.clone(),
                    });
                    illustration._keep_alive = Some(_keep_alive);
                }
                Ok(None) => {}
                Err(e) => log::error!("ugoira encode failed for {illust_id}: {e}"),
            }
        }
        Ok(illustration)
    }

    /// Fetches ugoira metadata (frame zip + frame delays) from the app API.
    pub async fn ugoira_metadata(&self, illust_id: u64) -> Result<UgoiraMetadataModel, PixivError> {
        let access_token = self.get_access_token().await?;
        let response = crate::site::CLIENT
            .get(format!(
                "{APP_API_URL}/v1/ugoira/metadata?illust_id={illust_id}"
            ))
            .header("app-os", "ios")
            .header("app-os-version", "14.6")
            .header("User-Agent", APP_USER_AGENT)
            .bearer_auth(access_token)
            .send()
            .await?;
        if !response.status().is_success() {
            return Err(PixivError::Api(format!("status {}", response.status())));
        }
        let json: serde_json::Value = serde_json::from_str(&response.text().await?)?;
        if json.get("error").is_some() {
            let message = json
                .get("message")
                .and_then(|v| v.as_str())
                .unwrap_or("ugoira metadata failed");
            return Err(PixivError::Api(message.to_string()));
        }
        let metadata = json
            .get("ugoira_metadata")
            .ok_or_else(|| PixivError::Api("missing ugoira_metadata".to_string()))?;
        Ok(serde_json::from_value(metadata.clone())?)
    }

    /// Downloads the frame zip and encodes one MP4 via ffmpeg. Returns the
    /// MP4 path plus the temp directory that must stay alive until the file
    /// is uploaded.
    async fn ugoira_video(
        &self,
        illust_id: u64,
    ) -> Result<Option<(String, tempfile::TempDir)>, PixivError> {
        if !crate::site::ffmpeg_available() {
            crate::site::log_once_ffmpeg_missing();
            return Ok(None);
        }
        let metadata = self.ugoira_metadata(illust_id).await?;
        if metadata.frames.is_empty() {
            return Ok(None);
        }
        let zip_url = metadata
            .zip_url
            .clone()
            .or_else(|| metadata.zip_urls.as_ref().map(|z| z.medium.clone()));
        let Some(zip_url) = zip_url else {
            return Ok(None);
        };
        // Stream the frame zip to a temp file instead of buffering it in
        // memory: ugoira zips can be hundreds of MB, and the old
        // download_media_limited path spiked RAM up to the size cap.
        let mut zip_file = tempfile::Builder::new()
            .suffix(".zip")
            .tempfile()
            .map_err(|e| PixivError::Api(format!("temp zip failed: {e}")))?;
        crate::site::download_media_to_file(&zip_url, 512 * 1024 * 1024, zip_file.as_file_mut())
            .await
            .map_err(|e| match e {
                FetchError::Http(e) => PixivError::Http(e),
                other => PixivError::Api(format!("frame zip download failed: {other}")),
            })?;
        let frame_delays = metadata.frames.iter().map(|f| f.delay).collect::<Vec<_>>();
        let result =
            tokio::task::spawn_blocking(move || -> Result<(String, tempfile::TempDir), String> {
                let frames_dir = tempfile::tempdir().map_err(|e| e.to_string())?;
                let out_dir = tempfile::tempdir().map_err(|e| e.to_string())?;

                // Extract frames to canonical zero-padded names; pixiv ugoira
                // frames are uniformly jpg or png per artwork. The zip is read
                // from disk; `zip_file` stays alive for the whole extraction.
                let mut archive = zip::ZipArchive::new(
                    std::fs::File::open(zip_file.path()).map_err(|e| e.to_string())?,
                )
                .map_err(|e| format!("unzip: {e}"))?;
                if archive.is_empty() {
                    return Err("empty frame zip".to_string());
                }
                // Uniform jpg or png per artwork; sniff the first entry's
                // magic bytes instead of trusting its filename.
                let first = archive.by_index(0).map_err(|e| e.to_string())?;
                let mut first_bytes = Vec::new();
                first
                    .take(64 * 1024 * 1024 + 1)
                    .read_to_end(&mut first_bytes)
                    .map_err(|e| e.to_string())?;
                if first_bytes.len() > 64 * 1024 * 1024 {
                    return Err("frame exceeds size cap".to_string());
                }
                let extension = if first_bytes.starts_with(&[0xFF, 0xD8]) {
                    "jpg"
                } else if first_bytes.starts_with(b"\x89PNG") {
                    "png"
                } else {
                    "jpg"
                };
                let mut count = 0usize;
                {
                    let path = frames_dir
                        .path()
                        .join(format!("img_{count:05}.{extension}"));
                    std::fs::write(&path, &first_bytes).map_err(|e| e.to_string())?;
                    count += 1;
                }
                for i in 1..archive.len() {
                    let entry = archive.by_index(i).map_err(|e| e.to_string())?;
                    if entry.size() > 64 * 1024 * 1024 {
                        return Err(format!("frame {i} exceeds size cap"));
                    }
                    let mut bytes = Vec::new();
                    entry
                        .take(64 * 1024 * 1024 + 1)
                        .read_to_end(&mut bytes)
                        .map_err(|e| e.to_string())?;
                    if bytes.len() > 64 * 1024 * 1024 {
                        return Err(format!("frame {i} exceeds size cap"));
                    }
                    let path = frames_dir
                        .path()
                        .join(format!("img_{count:05}.{extension}"));
                    std::fs::write(&path, bytes).map_err(|e| e.to_string())?;
                    count += 1;
                }
                if count == 0 {
                    return Err("empty frame zip".to_string());
                }

                // Constant rate from the median frame delay (ms).
                let mut delays = frame_delays;
                delays.sort_unstable();
                let median = delays[delays.len() / 2].max(1);
                let framerate = 1000.0 / median as f64;

                let output = out_dir.path().join("ugoira.mp4");
                let status = std::process::Command::new("ffmpeg")
                    .args([
                        "-y",
                        "-framerate",
                        &framerate.to_string(),
                        "-i",
                        &frames_dir
                            .path()
                            .join(format!("img_%05d.{extension}"))
                            .to_string_lossy(),
                        // libx264 needs even dimensions; pixiv ugoira frames can
                        // be odd-sized (e.g. 277x405).
                        "-vf",
                        "scale=trunc(iw/2)*2:trunc(ih/2)*2",
                        "-c:v",
                        "libx264",
                        "-pix_fmt",
                        "yuv420p",
                        "-movflags",
                        "+faststart",
                        &output.to_string_lossy(),
                    ])
                    .stdout(std::process::Stdio::null())
                    .stderr(std::process::Stdio::null())
                    .status()
                    .map_err(|e| format!("ffmpeg spawn failed: {e}"))?;
                if !status.success() {
                    return Err(format!("ffmpeg exited with {status}"));
                }
                Ok((output.to_string_lossy().into_owned(), out_dir))
            })
            .await
            .map_err(|e| {
                log::error!("ugoira encode worker panicked for {illust_id}: {e}");
                PixivError::Api(format!("ugoira worker failed: {e}"))
            })?;
        match result {
            Ok(pair) => Ok(Some(pair)),
            Err(message) => {
                log::error!("ugoira encode failed for {illust_id}: {message}");
                Ok(None)
            }
        }
    }
}

/// pixiv3-rs replacement: `None` when `PIXIV_REFRESH_TOKEN` is unset.
static PIXIV_CLIENT: LazyLock<Option<PixivAPI>> =
    LazyLock::new(|| env::var("PIXIV_REFRESH_TOKEN").ok().map(PixivAPI::new));

/// Set at startup when the login validation fails; pixiv stays disabled until
/// the next process start.
static DISABLED: AtomicBool = AtomicBool::new(false);

pub fn enabled() -> bool {
    !DISABLED.load(Ordering::Relaxed) && env::var("PIXIV_REFRESH_TOKEN").is_ok()
}

/// Permanently disables pixiv until the next process start.
pub fn disable() {
    DISABLED.store(true, Ordering::Relaxed);
}

/// Forces the refresh-token → access-token exchange now, surfacing invalid
/// tokens and network errors. Called once at bot startup; on failure the bot
/// calls [`disable`].
pub async fn validate() -> Result<(), PixivError> {
    match PIXIV_CLIENT.as_ref() {
        None => Err(PixivError::NoAuth),
        Some(client) => {
            client.get_access_token().await?;
            Ok(())
        }
    }
}

pub async fn fetch(illust_id: u64) -> Result<Illustration, FetchError> {
    let client = PIXIV_CLIENT
        .as_ref()
        .filter(|_| enabled())
        .ok_or(FetchError::Pixiv(PixivError::NoAuth))?;
    client.fetch(illust_id).await
}

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

    /// Skips when `PIXIV_REFRESH_TOKEN` is absent or empty (CI without the
    /// secret must stay green; GitHub Actions exposes an unset secret as an
    /// empty string, so `is_err()` alone is not enough).
    fn require_pixiv_token() -> bool {
        std::env::var("PIXIV_REFRESH_TOKEN")
            .ok()
            .filter(|s| !s.is_empty())
            .is_some()
    }

    #[tokio::test]
    async fn test_fetch() {
        dotenv().ok();
        if !require_pixiv_token() {
            eprintln!("skipping: no PIXIV_REFRESH_TOKEN");
            return;
        }
        let result = fetch(126839080).await;
        assert!(result.is_ok());
        println!("{:#?}", result);
    }

    #[tokio::test]
    #[ignore = "live network: requires outbound HTTPS to oauth.secure.pixiv.net"]
    async fn live_validate_with_bogus_token_fails() {
        dotenv().ok();
        // A bogus token must surface as Api error (invalid_grant), not panic.
        let client = PixivAPI::new("bogus_token_for_testing".to_string());
        let result = client.get_access_token().await;
        assert!(matches!(result, Err(PixivError::Api(_))), "got {result:?}");
    }
}