bunkr-client 0.2.2

#[cfg(feature = "download")]
use crate::TokenResponse;
use crate::core::types::{AlbumFile, FailedOperationInfo};
#[cfg(feature = "download")]
use crate::core::types::DownloadResponse;
use anyhow::{Result, anyhow};
#[cfg(feature = "download")]
use json5;
use std::sync::{Arc, Mutex};
#[cfg(feature = "download")]
use reqwest::{Client, header};
#[cfg(feature = "download")]
use regex::Regex;
#[cfg(feature = "download")]
use serde_json;
#[cfg(feature = "download")]
use std::path::Path;
#[cfg(feature = "download")]
use std::sync::OnceLock;
#[cfg(feature = "download")]
use std::time::Duration;
#[cfg(feature = "download")]
use tokio::io::AsyncWriteExt;
#[cfg(feature = "download")]
use tokio::task::JoinSet;

#[cfg(feature = "ui")]
use crate::ui::ui::UIState;
#[cfg(not(feature = "ui"))]
use crate::core::types::UIState;

#[cfg(not(feature = "ui"))]
impl UIState {
    pub fn add_current_operation(&mut self, _name: String, _progress: f64, _size: u64) {}
    pub fn update_progress(&mut self, _name: &str, _progress: f64) {}
    pub fn remove_current_operation(&mut self, _name: &str, _url: Option<&str>) {}
    pub fn add_failed_operation(&mut self, _name: String, _info: FailedOperationInfo) {}
    pub fn add_processed_bytes(&mut self, _bytes: u64) {}
}

pub struct BunkrDownloader {
    #[cfg(feature = "download")]
    client: Client,
    #[cfg(feature = "download")]
    headers: header::HeaderMap,
    #[cfg(feature = "download")]
    album_files_regex: OnceLock<Regex>,
    #[cfg(feature = "download")]
    trailing_comma_regex: OnceLock<Regex>,
    #[cfg(feature = "download")]
    keys_regex: OnceLock<Regex>,
    #[cfg(feature = "download")]
    id_regex: OnceLock<Regex>,
    #[cfg(feature = "download")]
    orig_regex: OnceLock<Regex>,
}

#[cfg(not(feature = "download"))]
impl BunkrDownloader {
    pub async fn new() -> Result<Self> {
        Err(anyhow!("Download feature is not enabled."))
    }

    pub async fn get_files(&self, _album_url: &str) -> Result<Vec<AlbumFile>> {
        Err(anyhow!("Download feature is not enabled."))
    }

    pub async fn download_file(&self, _file: &AlbumFile, _output_dir: &str, _ui_state: Option<Arc<Mutex<UIState>>>) -> Result<()> {
        Err(anyhow!("Download feature is not enabled."))
    }

    pub async fn download_files(&self, _files: Vec<AlbumFile>, _output_dir: &str, _batch_size: usize, _ui_state: Option<Arc<Mutex<UIState>>>) -> Result<()> {
        Err(anyhow!("Download feature is not enabled."))
    }
}

#[cfg(feature = "download")]
impl BunkrDownloader {
    pub async fn new() -> Result<Self> {
        let client = Client::new();

        let mut headers = header::HeaderMap::new();
        headers.insert("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36".parse()?);
        headers.insert("Accept", "application/json, text/plain, */*".parse()?);
        headers.insert("Accept-Language", "en-US,en;q=0.9".parse()?);
        headers.insert("Content-Type", "application/json".parse()?);
        headers.insert("Origin", "https://bunkr.su".parse()?);

        let album_files_regex = OnceLock::new();
        album_files_regex.get_or_init(|| Regex::new(r"window\.albumFiles\s*=\s*\[\s*([\s\S]*?)\s*\]\s*;").unwrap());

        let trailing_comma_regex = OnceLock::new();
        trailing_comma_regex.get_or_init(|| Regex::new(r",(\s*)([}\]])").unwrap());

        let keys_regex = OnceLock::new();
        keys_regex.get_or_init(|| Regex::new(r"(?m)^(\s*)(\w+)(\s*):").unwrap());

        let id_regex = OnceLock::new();
        id_regex.get_or_init(|| Regex::new(r#"data-file-id="(\d+)""#).unwrap());

        let orig_regex = OnceLock::new();
        orig_regex.get_or_init(|| Regex::new(r#"<h1 class="text-subs font-semibold text-base sm:text-lg truncate">([^<]+)</h1>"#).unwrap());

        Ok(Self {
            client,
            headers,
            album_files_regex,
            trailing_comma_regex,
            keys_regex,
            id_regex,
            orig_regex,
        })
    }

    pub async fn get_files(&self, album_url: &str) -> Result<Vec<AlbumFile>> {
        if album_url.contains("/a/") {
            self.get_album_files(album_url).await
        } else if album_url.contains("/f/") {
            let file = self.get_single_file(album_url).await?;
            Ok(vec![file])
        } else {
            Err(anyhow!("Unsupported URL: {}", album_url))
        }
    }

    async fn get_album_files(&self, album_url: &str) -> Result<Vec<AlbumFile>> {
        // Album URL
        // Ensure the URL has advanced=1
        let url = if album_url.contains("?") {
            format!("{}&advanced=1", album_url)
        } else {
            format!("{}?advanced=1", album_url)
        };

        let response = self.client.get(&url).send().await?;
        let html = response.text().await?;

        // Regex to extract the window.albumFiles array
        let re = self.album_files_regex.get().unwrap();

        let captures = re.captures(&html)
            .ok_or_else(|| anyhow!("Could not find albumFiles in the page"))?;

        let js_array = &captures[1];

        // Convert JavaScript object notation to JSON
        let json_str = self.js_to_json(js_array)?;

        // Parse the JS-like array using json5 (supports single quotes, trailing commas, etc.)
        let files: Vec<AlbumFile> = json5::from_str(&json_str)?;

        Ok(files)
    }

    async fn get_single_file(&self, file_url: &str) -> Result<AlbumFile> {
        // Individual file URL
        let response = self.client.get(file_url).send().await?;
        let html = response.text().await?;

        // Extract file id from <div id="fileTracker" data-file-id="...">
        let id_re = self.id_regex.get().unwrap();
        let id: i64 = id_re.captures(&html)
            .and_then(|c| c[1].parse().ok())
            .ok_or_else(|| anyhow!("Could not find file id"))?;

        // Extract original filename from <h1 class="text-subs font-semibold text-base sm:text-lg truncate">
        let orig_re = self.orig_regex.get().unwrap();
        let original = orig_re.captures(&html)
            .and_then(|c| Some(c[1].to_string()))
            .ok_or_else(|| anyhow!("Could not find file name"))?;

        // Create AlbumFile with extracted data, defaults for others
        let file = AlbumFile {
            id,
            name: original.clone(),
            original,
            slug: "".to_string(),
            file_type: "".to_string(),
            extension: "".to_string(),
            size: 0,
            timestamp: "".to_string(),
            thumbnail: "".to_string(),
            cdn_endpoint: "".to_string(),
        };

        Ok(file)
    }

    fn js_to_json(&self, js_str: &str) -> Result<String> {
        // Replace JavaScript object syntax with JSON
        let mut json = js_str.to_string();

        // Remove trailing commas before closing braces/brackets
        let re_trailing_comma = self.trailing_comma_regex.get().unwrap();
        json = re_trailing_comma.replace_all(&json, "$2").to_string();

        // Replace unquoted keys with quoted keys
        // This is a simple approach; for more complex cases, a proper JS parser might be needed
        let re_keys = self.keys_regex.get().unwrap();
        json = re_keys.replace_all(&json, "$1\"$2\"$3:").to_string();

        // Wrap in array brackets if not already
        let json = format!("[{}]", json);

        Ok(json)
    }

    pub async fn download_file(&self, file: &AlbumFile, output_dir: &str, ui_state: Option<Arc<Mutex<UIState>>>) -> Result<()> {
        Self::download_file_owned(
            self.client.clone(),
            self.headers.clone(),
            Self::owned_album_file(file),
            output_dir.to_string(),
            ui_state,
        ).await
    }

    fn owned_album_file(file: &AlbumFile) -> AlbumFile {
        AlbumFile {
            id: file.id,
            name: file.name.clone(),
            original: file.original.clone(),
            slug: file.slug.clone(),
            file_type: file.file_type.clone(),
            extension: file.extension.clone(),
            size: file.size,
            timestamp: file.timestamp.clone(),
            thumbnail: file.thumbnail.clone(),
            cdn_endpoint: file.cdn_endpoint.clone(),
        }
    }

    async fn fetch_json<T: serde::de::DeserializeOwned>(
        client: &Client,
        method: reqwest::Method,
        url: &str,
        headers: Option<&header::HeaderMap>,
        body: Option<serde_json::Value>,
        label: &str,
    ) -> Result<T> {
        let mut last_error = None;
        for attempt in 0..3 {
            if attempt > 0 {
                tokio::time::sleep(Duration::from_millis(500 * (1 << attempt))).await;
            }
            let mut builder = client.request(method.clone(), url);
            if let Some(h) = headers {
                builder = builder.headers(h.clone());
            }
            if let Some(ref b) = body {
                builder = builder.json(b);
            }
            match builder.send().await {
                Ok(resp) => {
                    let status = resp.status();
                    let body_text = resp.text().await?;
                    if !status.is_success() {
                        last_error = Some(anyhow!(
                            "{} API returned status {}: {}",
                            label, status, &body_text[..body_text.len().min(200)]
                        ));
                        continue;
                    }
                    match serde_json::from_str::<T>(&body_text) {
                        Ok(val) => return Ok(val),
                        Err(e) => {
                            last_error = Some(anyhow!(
                                "Failed to parse {} response (status {}): {} - body: {}",
                                label, status, e, &body_text[..body_text.len().min(200)]
                            ));
                            continue;
                        }
                    }
                }
                Err(e) => {
                    last_error = Some(anyhow!("{} request failed: {}", label, e));
                    continue;
                }
            }
        }
        Err(anyhow!(
            "{} failed after 3 retries: {}",
            label,
            last_error.unwrap()
        ))
    }

    async fn download_file_owned(
        client: Client,
        headers: header::HeaderMap,
        file: AlbumFile,
        output_dir: String,
        ui_state: Option<Arc<Mutex<UIState>>>,
    ) -> Result<()> {
        let file_path = Path::new(&output_dir).join(&file.original);
        if file_path.exists() {
            return Ok(());
        }

        let api_url = "https://dl.bunkr.cr/api/_001_v2";

        let mut last_error = None;
        for attempt in 0..3 {
            if attempt > 0 {
                tokio::time::sleep(Duration::from_millis(1000 * (1 << attempt))).await;
            }

            let body = serde_json::json!({ "id": file.id.to_string() });

            let download_resp = match Self::fetch_json::<DownloadResponse>(
                &client,
                reqwest::Method::POST,
                api_url,
                Some(&headers),
                Some(body),
                "download_url",
            ).await {
                Ok(r) => r,
                Err(e) => { last_error = Some(e); continue; }
            };

            let token_url = format!("https://glb-apisign.cdn.cr/sign?path={}", download_resp.path);

            let token_resp = match Self::fetch_json::<TokenResponse>(
                &client,
                reqwest::Method::GET,
                &token_url,
                None,
                None,
                "token",
            ).await {
                Ok(r) => r,
                Err(e) => { last_error = Some(e); continue; }
            };

            let full_url = format!("{}{}?n={}&token={}&ex={}", download_resp.mediafiles, download_resp.path, download_resp.original, token_resp.token, token_resp.ex);

            let mut download_headers = header::HeaderMap::new();
            download_headers.insert("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:151.0) Gecko/20100101 Firefox/151.0".parse()?);
            download_headers.insert("Accept", "*/*".parse()?);
            download_headers.insert("Accept-Language", "en-US,en;q=0.5".parse()?);
            download_headers.insert("Referer", "https://dl.bunkrr.cr/".parse()?);

            let response = match client.get(&full_url).headers(download_headers).send().await {
                Ok(r) => r,
                Err(e) => { last_error = Some(anyhow!("{}", e)); continue; }
            };

            if !response.status().is_success() {
                last_error = Some(anyhow!("Failed to download file: {}", response.status()));
                continue;
            }

            let total_size = response.content_length().unwrap_or(file.size as u64);
            let mut downloaded = 0u64;

            match tokio::fs::File::create(&file_path).await {
                Ok(mut file_handle) => {
                    let result: Result<()> = async {
                        let mut stream = response;
                        while let Some(chunk) = stream.chunk().await.map_err(|e| anyhow!("{}", e))? {
                            file_handle.write_all(&chunk).await?;
                            downloaded += chunk.len() as u64;

                            if let Some(ref state) = ui_state {
                                let mut state = state.lock().unwrap();
                                let progress = if total_size > 0 { (downloaded as f64 / total_size as f64).min(1.0) } else { 0.0 };
                                state.update_progress(&file.original, progress);
                                state.add_processed_bytes(chunk.len() as u64);
                            }
                        }
                        Ok(())
                    }.await;

                    match result {
                        Ok(()) => return Ok(()),
                        Err(e) => {
                            let _ = tokio::fs::remove_file(&file_path).await;
                            last_error = Some(e);
                            continue;
                        }
                    }
                }
                Err(e) => {
                    last_error = Some(anyhow!("Failed to create file: {}", e));
                    continue;
                }
            }
        }

        Err(anyhow!(
            "Download failed after 3 retries: {}",
            last_error.unwrap()
        ))
    }

    pub async fn download_files(&self, files: Vec<AlbumFile>, output_dir: &str, batch_size: usize, ui_state: Option<Arc<Mutex<UIState>>>) -> Result<()> {
        let batch_size = batch_size.max(1);
        let output_dir = output_dir.to_string();
        let client = self.client.clone();
        let headers = self.headers.clone();

        let mut files_iter = files.into_iter();
        let mut join_set = JoinSet::new();

        let mut spawn_next = |join_set: &mut JoinSet<(AlbumFile, Result<()>)>| {
            let Some(file) = files_iter.next() else {
                return;
            };

            let client = client.clone();
            let headers = headers.clone();
            let output_dir = output_dir.clone();
            let ui_state = ui_state.clone();

            join_set.spawn(async move {
                if let Some(ref state) = ui_state {
                    let mut state = state.lock().unwrap();
                    state.add_current_operation(file.original.clone(), 0.0, file.size as u64);
                }

                let file_for_result = BunkrDownloader::owned_album_file(&file);
                let result = BunkrDownloader::download_file_owned(client, headers, file, output_dir, ui_state.clone()).await;
                (file_for_result, result)
            });
        };

        for _ in 0..batch_size {
            spawn_next(&mut join_set);
        }

        while let Some(result) = join_set.join_next().await {
            if let Ok((file, result)) = result {
                match result {
                    Ok(_) => {
                        if let Some(ref state) = ui_state {
                            let mut state = state.lock().unwrap();
                            state.remove_current_operation(&file.original, None);
                        }
                    }
                    Err(e) => {
                        if let Some(ref state) = ui_state {
                            let mut state = state.lock().unwrap();
                            let info = FailedOperationInfo {
                                path: file.original.clone(),
                                error: e.to_string(),
                                file_size: file.size as u64,
                                status_code: None,
                            };
                            state.add_failed_operation(file.original.clone(), info);
                        }
                    }
                }
            }

            spawn_next(&mut join_set);
        }

        Ok(())
    }
}