jellytau/src-tauri/src/download/worker.rs

//! Download worker for HTTP streaming with progress tracking and retry logic

use log::warn;
use std::time::Duration;

use futures_util::StreamExt;
use tokio::fs;
use tokio::io::AsyncWriteExt;

use super::DownloadTask;

/// Download worker that handles individual download tasks
pub struct DownloadWorker {
    /// HTTP client for downloads
    client: reqwest::Client,
    /// Maximum retry attempts
    max_retries: u32,
}

impl DownloadWorker {
    pub fn new() -> Self {
        let client = reqwest::Client::builder()
            .timeout(Duration::from_secs(300)) // 5 minute timeout
            .https_only(true)
            .build()
            .expect("Failed to create HTTP client");

        Self {
            client,
            max_retries: 3,
        }
    }

    /// Download a file with retry logic and progress tracking
    pub async fn download<F>(
        &self,
        task: &DownloadTask,
        on_progress: F,
    ) -> Result<DownloadResult, DownloadError>
    where
        F: Fn(u64, Option<u64>) + Send + Sync,
    {
        let mut retries = 0;

        loop {
            match self.try_download(task, &on_progress).await {
                Ok(result) => return Ok(result),
                Err(e) if retries < self.max_retries && e.is_retryable() => {
                    retries += 1;
                    let delay = Self::exponential_backoff(retries);
                    warn!(
                        "Download failed (attempt {}/{}), retrying in {:?}: {}",
                        retries, self.max_retries, delay, e
                    );
                    tokio::time::sleep(delay).await;
                }
                Err(e) => return Err(e),
            }
        }
    }

    /// Attempt a single download
    async fn try_download<F>(&self, task: &DownloadTask, on_progress: &F) -> Result<DownloadResult, DownloadError>
    where
        F: Fn(u64, Option<u64>) + Send + Sync,
    {
        // Create parent directories
        if let Some(parent) = task.target_path.parent() {
            fs::create_dir_all(parent)
                .await
                .map_err(|e| DownloadError::FileSystem(e.to_string()))?;
        }

        // Check for partial download
        let temp_path = task.target_path.with_extension("part");
        let existing_bytes = if temp_path.exists() {
            fs::metadata(&temp_path)
                .await
                .map(|m| m.len())
                .unwrap_or(0)
        } else {
            0
        };

        // Build HTTP request with Range header for resume support
        let mut request = self.client.get(&task.url);
        if existing_bytes > 0 {
            request = request.header("Range", format!("bytes={}-", existing_bytes));
        }

        // Send request
        let response = request
            .send()
            .await
            .map_err(|e| DownloadError::Network(e.to_string()))?;

        // Check status
        if !response.status().is_success() && response.status().as_u16() != 206 {
            return Err(DownloadError::Http(response.status().as_u16()));
        }

        // Get content length
        let _total_bytes = response
            .headers()
            .get(reqwest::header::CONTENT_LENGTH)
            .and_then(|v| v.to_str().ok())
            .and_then(|v| v.parse::<u64>().ok())
            .map(|len| if existing_bytes > 0 { len + existing_bytes } else { len });

        // Open file for appending
        let mut file = if existing_bytes > 0 {
            fs::OpenOptions::new()
                .append(true)
                .open(&temp_path)
                .await
        } else {
            fs::File::create(&temp_path).await
        }
        .map_err(|e| DownloadError::FileSystem(e.to_string()))?;

        // Stream download with progress tracking
        let mut downloaded = existing_bytes;
        let mut stream = response.bytes_stream();
        let mut last_progress_emit = std::time::Instant::now();

        while let Some(chunk) = stream.next().await {
            let chunk = chunk.map_err(|e| DownloadError::Network(e.to_string()))?;

            file.write_all(&chunk)
                .await
                .map_err(|e| DownloadError::FileSystem(e.to_string()))?;

            downloaded += chunk.len() as u64;

            // Emit progress every 500ms or every MB
            if last_progress_emit.elapsed() > Duration::from_millis(500)
                || downloaded % (1024 * 1024) == 0
            {
                last_progress_emit = std::time::Instant::now();
                on_progress(downloaded, _total_bytes);
            }
        }

        file.sync_all()
            .await
            .map_err(|e| DownloadError::FileSystem(e.to_string()))?;

        // Move from .part to final location
        fs::rename(&temp_path, &task.target_path)
            .await
            .map_err(|e| DownloadError::FileSystem(e.to_string()))?;

        Ok(DownloadResult {
            bytes_downloaded: downloaded,
        })
    }

    /// Calculate exponential backoff delay
    fn exponential_backoff(retry_count: u32) -> Duration {
        let base_delay = 5; // 5 seconds
        let delay_secs = base_delay * 3u64.pow(retry_count - 1); // 5s, 15s, 45s
        Duration::from_secs(delay_secs)
    }
}

/// Result of a successful download
#[derive(Debug)]
pub struct DownloadResult {
    pub bytes_downloaded: u64,
}

/// Download error types
#[derive(Debug)]
pub enum DownloadError {
    Network(String),
    Http(u16),
    FileSystem(String),
}

impl DownloadError {
    /// Check if this error is retryable
    fn is_retryable(&self) -> bool {
        match self {
            DownloadError::Network(_) => true,
            DownloadError::Http(status) => *status >= 500, // Retry server errors
            DownloadError::FileSystem(_) => false,
        }
    }
}

impl std::fmt::Display for DownloadError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            DownloadError::Network(msg) => write!(f, "Network error: {}", msg),
            DownloadError::Http(status) => write!(f, "HTTP error {}", status),
            DownloadError::FileSystem(msg) => write!(f, "File system error: {}", msg),
        }
    }
}

impl std::error::Error for DownloadError {}

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

    #[test]
    fn test_exponential_backoff() {
        assert_eq!(DownloadWorker::exponential_backoff(1), Duration::from_secs(5));
        assert_eq!(DownloadWorker::exponential_backoff(2), Duration::from_secs(15));
        assert_eq!(DownloadWorker::exponential_backoff(3), Duration::from_secs(45));
    }

    #[test]
    fn test_error_retryable() {
        assert!(DownloadError::Network("timeout".to_string()).is_retryable());
        assert!(DownloadError::Http(500).is_retryable());
        assert!(DownloadError::Http(503).is_retryable());
        assert!(!DownloadError::Http(404).is_retryable());
        assert!(!DownloadError::FileSystem("disk full".to_string()).is_retryable());
    }
}