oma-fetch 0.40.0

use std::{borrow::Cow, cmp::Ordering, fmt::Debug, path::PathBuf, sync::Arc, time::Duration};

use ahash::AHashMap;
use bon::Builder;
use checksum::Checksum;
use download::{BuilderError, SingleDownloader, SuccessSummary};

use oma_logger::debug;
use reqwest::{Method, Response, Url};
use reqwest_middleware::{ClientWithMiddleware, RequestBuilder};
use serde::{Deserialize, Serialize};
use tokio::task::JoinSet;

pub mod checksum;
pub mod download;
pub use crate::download::SingleDownloadError;

pub use reqwest;

#[derive(Clone, Default, Builder)]
pub struct DownloadEntry {
    pub source: Vec<DownloadSource>,
    pub filename: String,
    dir: PathBuf,
    final_dir: Option<PathBuf>,
    hash: Option<Checksum>,
    allow_resume: bool,
    msg: Option<Cow<'static, str>>,
    /// Whether this entry downloads the by-hash variant of an index file with
    /// a fallback to the traditional by-name path of the same file. When
    /// `true`, a `NextUrl` event means a same-file fallback rather than a
    /// different mirror.
    #[builder(default)]
    by_hash_fallback: bool,
}

impl Debug for DownloadEntry {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DownloadEntry")
            .field("source", &self.source)
            .field("filename", &self.filename)
            .field("dir", &self.dir)
            .field("hash", &self.hash.as_ref().map(|c| c.to_string()))
            .field("allow_resume", &self.allow_resume)
            .field("msg", &self.msg)
            .field("by_hash_fallback", &self.by_hash_fallback)
            .finish()
    }
}

#[derive(Debug, Clone, Default, PartialEq, Eq, Copy)]
pub enum CompressType {
    Bz2,
    Gzip,
    Xz,
    Zstd,
    Lzma,
    Lz4,
    #[default]
    None,
}

impl From<&str> for CompressType {
    fn from(s: &str) -> Self {
        match s {
            "xz" => CompressType::Xz,
            "gz" => CompressType::Gzip,
            "bz2" => CompressType::Bz2,
            "zst" => CompressType::Zstd,
            _ => CompressType::None,
        }
    }
}

#[derive(Debug, Clone)]
pub struct DownloadSource {
    pub url: String,
    pub source_type: DownloadSourceType,
    /// The compression format of the file this source points to. Used to
    /// decompress the downloaded file. This is per-source so that a single
    /// [`DownloadEntry`] can fall back between different compression formats
    /// (e.g. `Packages.zst` → `Packages.xz` → `Packages.gz`), matching apt's
    /// `Acquire::CompressionTypes::Order` behavior.
    pub file_type: CompressType,
}

#[derive(Debug, PartialEq, Eq, Clone)]
pub enum DownloadSourceType {
    Http,
    Local(bool),
}

impl PartialOrd for DownloadSourceType {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for DownloadSourceType {
    fn cmp(&self, other: &Self) -> Ordering {
        match self {
            DownloadSourceType::Http => match other {
                DownloadSourceType::Http => Ordering::Equal,
                DownloadSourceType::Local { .. } => Ordering::Less,
            },
            DownloadSourceType::Local { .. } => match other {
                DownloadSourceType::Http => Ordering::Greater,
                DownloadSourceType::Local { .. } => Ordering::Equal,
            },
        }
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub enum Event {
    ChecksumMismatch {
        index: usize,
        filename: String,
        times: usize,
    },
    Timeout {
        filename: String,
        times: usize,
    },
    GlobalProgressAdd(u64),
    GlobalProgressSub(u64),
    ProgressDone(usize),
    NewProgressSpinner {
        index: usize,
        total: usize,
        msg: String,
    },
    NewProgressBar {
        index: usize,
        total: usize,
        msg: String,
        size: u64,
    },
    ProgressInc {
        index: usize,
        size: u64,
    },
    NextUrl {
        index: usize,
        file_name: String,
        err: SingleDownloadError,
        /// `true` when the failed source was the by-hash variant and the next
        /// source is the traditional by-name path of the same file.
        by_hash: bool,
    },
    DownloadDone {
        index: usize,
        msg: Box<str>,
    },
    Failed {
        file_name: String,
        error: SingleDownloadError,
    },
    AllDone,
    NewGlobalProgressBar(u64),
}

#[derive(Serialize, Deserialize)]
pub(crate) enum SingleDownloadErrorHelper {
    SetPermission { source: String },
    OpenAsWriteMode { source: String },
    Open { source: String },
    Create { source: String },
    Seek { source: String },
    Write { source: String },
    Flush { source: String },
    Remove { source: String },
    CreateSymlink { source: String },
    ReqwestMiddlewareError { source: String },
    BrokenPipe { source: String },
    SendRequestTimeout,
    DownloadTimeout,
    ChecksumMismatch,
    AcquireError,
}

#[derive(Builder)]
pub struct DownloadManager {
    client: ClientWithMiddleware,
    download_list: Box<[DownloadEntry]>,
    #[builder(default = 4)]
    threads: usize,
    #[builder(default = 3)]
    retry_times: usize,
    #[builder(default)]
    total_size: u64,
    #[builder(default = Duration::from_secs(15))]
    timeout: Duration,
}

#[derive(Debug)]
pub struct Summary {
    pub success: Vec<SuccessSummary>,
    pub failed: Vec<String>,
}

impl Summary {
    pub fn is_download_success(&self) -> bool {
        self.failed.is_empty()
    }

    pub fn has_wrote(&self) -> bool {
        self.success.iter().any(|x| x.wrote)
    }
}

impl DownloadManager {
    /// Start download
    pub async fn start_download<F, Fut>(mut self, callback: F) -> Result<Summary, BuilderError>
    where
        F: Fn(Event) -> Fut + Send + Sync + 'static,
        Fut: std::future::Future<Output = ()> + Send + 'static,
    {
        let mut list = vec![];
        let len = self.download_list.len();

        let mut source_locks = AHashMap::new();

        for (i, c) in std::mem::take(&mut self.download_list)
            .into_iter()
            .enumerate()
        {
            let source_key = if let Some(src) = c.source.first() {
                if let Ok(url) = Url::parse(&src.url) {
                    format!("{}://{}", url.scheme(), url.host_str().unwrap_or("unknown"))
                } else {
                    "fallback_source".to_string()
                }
            } else {
                "unknown_source".to_string()
            };

            let source_sem = source_locks
                .entry(source_key)
                .or_insert_with(|| Arc::new(tokio::sync::Semaphore::new(self.threads)))
                .clone();

            let single = SingleDownloader::builder()
                .client(self.client.clone())
                .download_list_index(i)
                .entry(c)
                .total(len)
                .retry_times(self.retry_times)
                .timeout(self.timeout)
                .build()?;

            list.push((single, source_sem));
        }

        if self.total_size != 0 {
            callback(Event::NewGlobalProgressBar(self.total_size)).await;
        }

        let mut set = JoinSet::new();
        let callback_arc = Arc::new(callback);

        for (single, source_sem) in list {
            let cb = callback_arc.clone();

            set.spawn(async move {
                let _permit = match source_sem.acquire_owned().await {
                    Ok(p) => Some(p),
                    Err(_) => {
                        cb(Event::Failed {
                            file_name: single.entry.filename.to_string(),
                            error: SingleDownloadError::AcquireError,
                        })
                        .await;

                        return download::DownloadResult::Failed {
                            file_name: single.entry.filename,
                        };
                    }
                };

                single.try_download(&*cb).await
            });
        }

        let mut success = vec![];
        let mut failed = vec![];

        while let Some(res) = set.join_next().await {
            match res {
                Ok(download::DownloadResult::Success(success_summary)) => {
                    success.push(success_summary);
                }
                Ok(download::DownloadResult::Failed { file_name }) => {
                    failed.push(file_name);
                }
                Err(_) => {
                    failed.push("task_panicked".to_string());
                }
            }
        }

        callback_arc(Event::AllDone).await;

        Ok(Summary { success, failed })
    }
}

pub async fn send_request_with_url_and_method(
    url: &str,
    client: &ClientWithMiddleware,
    method: Method,
) -> Result<Response, reqwest_middleware::Error> {
    let resp = client.request(method, url).send().await?;
    let headers = resp.headers();

    debug!(
        "\nDownload URL: {url}\nStatus: {}\nHeaders: {headers:#?}",
        resp.status()
    );

    let resp = resp.error_for_status()?;

    Ok(resp)
}

pub async fn send_request(request: RequestBuilder) -> Result<Response, reqwest_middleware::Error> {
    let resp = request.send().await?;
    let headers = resp.headers();
    let url = resp.url();

    debug!(
        "\nDownload URL: {url}\nStatus: {}\nHeaders: {headers:#?}",
        resp.status()
    );

    let resp = resp.error_for_status()?;

    Ok(resp)
}