dataforge/

lib.rs

//! Fast dataset discovery and archive helpers for `thesa` and other Rust tools.
//!
//! Made by Trevor Knott for Knott Dynamics.
//!
//! DataForge normalizes dataset targets, queries public dataset catalogs, and
//! archives dataset files with bounded concurrent downloads. The first-class
//! provider is Hugging Face Datasets; Zenodo is included for DOI/record-backed
//! research datasets.
//!
//! ```
//! assert_eq!(
//!     dataforge::parse_dataset_target_for_provider(
//!         "https://huggingface.co/datasets/squad",
//!         dataforge::DatasetProvider::Hf,
//!     )
//!     .unwrap(),
//!     dataforge::DatasetTarget::Dataset("squad".to_string())
//! );
//! assert_eq!(
//!     dataforge::DatasetProvider::from_slug("hugging-face"),
//!     Some(dataforge::DatasetProvider::Hf)
//! );
//! ```

use std::collections::{BTreeSet, VecDeque};
use std::fmt;
use std::fs::{self, File};
use std::io;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::thread;

use reqwest::StatusCode;
use reqwest::blocking::{Client, Response};
use serde::{Deserialize, Serialize};

const HF_BASE: &str = "https://huggingface.co";
const ZENODO_API: &str = "https://zenodo.org/api";
const DEFAULT_PAGE_SIZE: usize = 100;
const DEFAULT_CONCURRENCY: usize = 4;
const DEFAULT_USER_AGENT: &str = "dataforge/0.1";

/// Result type used by DataForge APIs.
pub type Result<T> = std::result::Result<T, DataForgeError>;

/// Supported dataset providers.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum DatasetProvider {
    /// Hugging Face Datasets on the Hugging Face Hub.
    Hf,
    /// Zenodo records and files.
    Zenodo,
}

impl DatasetProvider {
    /// Stable provider list in CLI display order.
    pub const ALL: [Self; 2] = [Self::Hf, Self::Zenodo];

    /// Return all supported providers in stable CLI display order.
    pub fn all() -> &'static [Self] {
        &Self::ALL
    }

    /// Human-readable provider label.
    pub fn label(self) -> &'static str {
        match self {
            Self::Hf => "Hugging Face Datasets",
            Self::Zenodo => "Zenodo",
        }
    }

    /// Compact UI label.
    pub fn short_label(self) -> &'static str {
        match self {
            Self::Hf => "HF",
            Self::Zenodo => "ZEN",
        }
    }

    /// Stable lowercase provider slug.
    pub fn slug(self) -> &'static str {
        match self {
            Self::Hf => "hf",
            Self::Zenodo => "zenodo",
        }
    }

    /// Parse common provider aliases accepted by thesa-style CLIs.
    pub fn from_slug(value: &str) -> Option<Self> {
        let normalized = normalized_slug(value);
        match normalized.as_str() {
            "hf" | "huggingface" | "huggingfacedatasets" | "hfdatasets" => Some(Self::Hf),
            "zen" | "zenodo" => Some(Self::Zenodo),
            _ => None,
        }
    }
}

impl fmt::Display for DatasetProvider {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.slug())
    }
}

/// Normalized dataset target requested by a user or CLI.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum DatasetTarget {
    /// Provider namespace, owner, or organization.
    Namespace(String),
    /// Concrete provider dataset id or record id.
    Dataset(String),
    /// Provider search query.
    Search(String),
    /// Provider-specific top/trending/popular preset.
    Top,
    /// Provider-specific latest/newest preset.
    Latest,
}

impl DatasetTarget {
    /// Stable target kind label.
    pub fn kind(&self) -> &'static str {
        match self {
            Self::Namespace(_) => "namespace",
            Self::Dataset(_) => "dataset",
            Self::Search(_) => "search",
            Self::Top => "top",
            Self::Latest => "latest",
        }
    }

    /// Borrow the contained value for namespace, dataset, or search targets.
    pub fn value(&self) -> Option<&str> {
        match self {
            Self::Namespace(value) | Self::Dataset(value) | Self::Search(value) => Some(value),
            Self::Top | Self::Latest => None,
        }
    }

    /// Return true when this target points at one concrete dataset/record.
    pub fn is_dataset(&self) -> bool {
        matches!(self, Self::Dataset(_))
    }

    /// Return true when this target may expand into many datasets.
    pub fn is_collection(&self) -> bool {
        matches!(
            self,
            Self::Namespace(_) | Self::Search(_) | Self::Top | Self::Latest
        )
    }
}

/// Lightweight dataset metadata returned by discovery calls.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct DatasetSummary {
    /// Provider that produced the dataset.
    pub provider: DatasetProvider,
    /// Stable provider id, such as `owner/name` or a Zenodo record id.
    pub id: String,
    /// Human-readable title when the provider exposes one.
    pub title: Option<String>,
    /// Download count or similar popularity signal.
    pub downloads: Option<u64>,
    /// Like/favorite count when available.
    pub likes: Option<u64>,
    /// Last modification timestamp as returned by the provider.
    pub last_modified: Option<String>,
    /// Short description when available.
    pub description: Option<String>,
    /// Provider tags, keywords, or topics.
    pub tags: Vec<String>,
    /// Number of downloadable files if known.
    pub file_count: Option<usize>,
    /// Total known file size in bytes.
    pub size: Option<u64>,
    /// Canonical web URL for humans.
    pub url: Option<String>,
}

/// One downloadable file in a dataset archive.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct DatasetFile {
    /// Provider-relative file path.
    pub path: String,
    /// Direct download URL.
    pub download_url: String,
    /// File size in bytes when known.
    pub size: Option<u64>,
    /// Provider checksum string when known.
    pub checksum: Option<String>,
}

/// Full dataset metadata plus files to archive.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct DatasetDetails {
    /// Dataset summary.
    pub summary: DatasetSummary,
    /// Downloadable files.
    pub files: Vec<DatasetFile>,
}

/// Options controlling archive writes.
#[derive(Debug, Clone)]
pub struct ArchiveOptions {
    /// Directory where dataset subdirectories are written.
    pub output: PathBuf,
    /// Maximum number of files downloaded at once.
    pub concurrency: usize,
    /// Skip a dataset directory when it already exists.
    pub skip_existing: bool,
    /// Optional substring filter applied to discovered dataset ids and titles.
    pub filter: Option<String>,
}

impl ArchiveOptions {
    /// Create archive options for an output directory.
    pub fn new(output: impl Into<PathBuf>) -> Self {
        Self {
            output: output.into(),
            ..Self::default()
        }
    }
}

impl Default for ArchiveOptions {
    fn default() -> Self {
        Self {
            output: PathBuf::from("datasets"),
            concurrency: DEFAULT_CONCURRENCY,
            skip_existing: false,
            filter: None,
        }
    }
}

/// One archived or skipped file.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ArchivedFile {
    /// Dataset id that owns the file.
    pub dataset_id: String,
    /// Provider-relative file path.
    pub source_path: String,
    /// Local path written or skipped.
    pub local_path: PathBuf,
    /// Number of bytes written during this run.
    pub bytes_written: u64,
    /// Whether an existing file was skipped.
    pub skipped: bool,
}

/// Failed dataset or file archive entry.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ArchiveFailure {
    /// Dataset id that failed.
    pub dataset_id: String,
    /// Optional provider-relative file path.
    pub file: Option<String>,
    /// Human-readable failure message.
    pub message: String,
}

/// Summary returned after an archive run.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ArchiveSummary {
    /// Number of datasets discovered before filtering and archiving.
    pub discovered: usize,
    /// Number of datasets selected after filter application.
    pub selected: usize,
    /// Number of dataset directories completed without file failures.
    pub archived: usize,
    /// Number of dataset directories skipped because they already existed.
    pub skipped: usize,
    /// Number of files written.
    pub files_archived: usize,
    /// Bytes written during this run.
    pub bytes_written: u64,
    /// Per-file successful writes or skips.
    pub files: Vec<ArchivedFile>,
    /// Dataset or file failures.
    pub failed: Vec<ArchiveFailure>,
}

impl ArchiveSummary {
    /// Return true when no archive failures occurred.
    pub fn is_success(&self) -> bool {
        self.failed.is_empty()
    }

    /// Total attempted datasets after filtering.
    pub fn attempted(&self) -> usize {
        self.archived + self.skipped + self.failed_dataset_count()
    }

    /// Number of unique datasets with at least one failure.
    pub fn failed_dataset_count(&self) -> usize {
        self.failed
            .iter()
            .map(|failure| failure.dataset_id.as_str())
            .collect::<BTreeSet<_>>()
            .len()
    }
}

/// Errors returned by DataForge.
#[derive(Debug, thiserror::Error)]
pub enum DataForgeError {
    /// The provided target string could not be normalized.
    #[error("invalid dataset target '{0}'")]
    InvalidTarget(String),

    /// The requested provider target is not supported by the current operation.
    #[error("unsupported dataset provider target: {0}")]
    UnsupportedTarget(String),

    /// Provider API returned a non-success status.
    #[error("{provider} API returned {status}: {body}")]
    ProviderApi {
        /// Provider slug.
        provider: &'static str,
        /// HTTP status code.
        status: u16,
        /// Response body or fallback message.
        body: String,
    },

    /// Provider did not find the requested dataset or record.
    #[error("{provider} dataset '{target}' was not found")]
    NotFound {
        /// Provider slug.
        provider: &'static str,
        /// Requested target.
        target: String,
    },

    /// Download failed for one dataset file.
    #[error("download failed for '{dataset}' file '{file}': {message}")]
    DownloadFailed {
        /// Dataset id.
        dataset: String,
        /// Provider-relative file path.
        file: String,
        /// Human-readable failure message.
        message: String,
    },

    /// Output path exists but is not a directory.
    #[error("output path '{}' is not a directory", .0.display())]
    OutputNotDirectory(PathBuf),

    /// Filesystem I/O failure.
    #[error("I/O error: {0}")]
    Io(#[from] io::Error),

    /// HTTP request failure.
    #[error("request error: {0}")]
    Request(#[from] reqwest::Error),

    /// JSON serialization failure.
    #[error("JSON error: {0}")]
    Json(#[from] serde_json::Error),
}

/// Blocking dataset discovery and archive client.
#[derive(Clone)]
pub struct DataForge {
    provider: DatasetProvider,
    client: Client,
    token: Option<String>,
    page_size: usize,
}

impl DataForge {
    /// Build a new client for the selected provider.
    pub fn new(provider: DatasetProvider) -> Result<Self> {
        Self::with_user_agent(provider, DEFAULT_USER_AGENT)
    }

    /// Build a new client with a custom User-Agent.
    pub fn with_user_agent(provider: DatasetProvider, user_agent: &str) -> Result<Self> {
        let client = Client::builder().user_agent(user_agent).build()?;
        Ok(Self {
            provider,
            client,
            token: None,
            page_size: DEFAULT_PAGE_SIZE,
        })
    }

    /// Attach a provider token. Hugging Face uses `Authorization: Bearer <token>`.
    pub fn with_token(mut self, token: impl Into<String>) -> Self {
        let token = token.into();
        if !token.trim().is_empty() {
            self.token = Some(token);
        }
        self
    }

    /// Set the API page size used for list requests.
    pub fn with_page_size(mut self, page_size: usize) -> Self {
        self.page_size = page_size.max(1);
        self
    }

    /// Return the provider this client targets.
    pub fn provider(&self) -> DatasetProvider {
        self.provider
    }

    /// Discover datasets for a normalized target.
    pub fn discover(&self, target: &DatasetTarget) -> Result<Vec<DatasetSummary>> {
        match self.provider {
            DatasetProvider::Hf => {
                discover_hf(&self.client, target, self.token.as_deref(), self.page_size)
            }
            DatasetProvider::Zenodo => discover_zenodo(&self.client, target, self.page_size),
        }
    }

    /// Fetch full details and file URLs for one dataset id.
    pub fn inspect(&self, dataset_id: &str) -> Result<DatasetDetails> {
        match self.provider {
            DatasetProvider::Hf => inspect_hf(&self.client, dataset_id, self.token.as_deref()),
            DatasetProvider::Zenodo => inspect_zenodo(&self.client, dataset_id),
        }
    }

    /// Discover and archive datasets for a target.
    pub fn archive_target(
        &self,
        target: &DatasetTarget,
        options: &ArchiveOptions,
    ) -> Result<ArchiveSummary> {
        let datasets = self.discover(target)?;
        self.archive_datasets(&datasets, options)
    }

    /// Archive discovered dataset summaries with bounded concurrent file downloads.
    pub fn archive_datasets(
        &self,
        datasets: &[DatasetSummary],
        options: &ArchiveOptions,
    ) -> Result<ArchiveSummary> {
        prepare_output(&options.output)?;

        let mut summary = ArchiveSummary {
            discovered: datasets.len(),
            ..ArchiveSummary::default()
        };
        let selected = datasets
            .iter()
            .filter(|dataset| dataset_matches_filter(dataset, options.filter.as_deref()))
            .cloned()
            .collect::<Vec<_>>();
        summary.selected = selected.len();

        let mut jobs = VecDeque::new();
        let mut scheduled_datasets = BTreeSet::new();

        for dataset in selected {
            let dataset_dir = options.output.join(dataset_dir_name(&dataset.id));
            if dataset_dir.exists() {
                if options.skip_existing {
                    summary.skipped += 1;
                    continue;
                }
                if dataset_dir.is_dir() {
                    fs::remove_dir_all(&dataset_dir)?;
                } else {
                    fs::remove_file(&dataset_dir)?;
                }
            }
            fs::create_dir_all(&dataset_dir)?;

            let details = match self.inspect(&dataset.id) {
                Ok(details) => details,
                Err(error) => {
                    summary.failed.push(ArchiveFailure {
                        dataset_id: dataset.id,
                        file: None,
                        message: error.to_string(),
                    });
                    continue;
                }
            };

            write_dataset_manifest(&dataset_dir, &details)?;
            scheduled_datasets.insert(details.summary.id.clone());

            for file in details.files {
                let local_path = dataset_dir.join(safe_relative_path(&file.path));
                jobs.push_back(DownloadJob {
                    dataset_id: details.summary.id.clone(),
                    source_path: file.path,
                    url: file.download_url,
                    local_path,
                });
            }
        }

        if jobs.is_empty() {
            summary.archived += scheduled_datasets.len();
            return Ok(summary);
        }

        let results = run_download_jobs(
            self.client.clone(),
            self.token.clone(),
            options.concurrency,
            jobs,
        );
        let mut failed_datasets = BTreeSet::new();

        for result in results {
            match result {
                Ok(file) => {
                    summary.bytes_written += file.bytes_written;
                    if !file.skipped {
                        summary.files_archived += 1;
                    }
                    summary.files.push(file);
                }
                Err(failure) => {
                    failed_datasets.insert(failure.dataset_id.clone());
                    summary.failed.push(failure);
                }
            }
        }

        summary.archived += scheduled_datasets.difference(&failed_datasets).count();

        Ok(summary)
    }
}

/// Parse a Hugging Face-style dataset target.
pub fn parse_dataset_target(input: &str) -> Result<DatasetTarget> {
    parse_dataset_target_for_provider(input, DatasetProvider::Hf)
}

/// Parse a dataset target using provider-specific rules.
pub fn parse_dataset_target_for_provider(
    input: &str,
    provider: DatasetProvider,
) -> Result<DatasetTarget> {
    let normalized = input.trim();
    if normalized.is_empty() {
        return Err(DataForgeError::InvalidTarget(
            "empty dataset target".to_string(),
        ));
    }

    let canonical = normalized.to_lowercase();
    match canonical.as_str() {
        "top" | "trending" | "popular" => return Ok(DatasetTarget::Top),
        "latest" | "newest" | "new" => return Ok(DatasetTarget::Latest),
        _ => {}
    }

    if let Some(value) = normalized.strip_prefix("dataset:") {
        return non_empty_target(value, DatasetTarget::Dataset);
    }
    if let Some(value) = normalized.strip_prefix("search:") {
        return non_empty_target(value, DatasetTarget::Search);
    }
    if let Some(value) = normalized
        .strip_prefix("org:")
        .or_else(|| normalized.strip_prefix("owner:"))
        .or_else(|| normalized.strip_prefix("namespace:"))
    {
        return non_empty_target(value, DatasetTarget::Namespace);
    }

    match provider {
        DatasetProvider::Hf => parse_hf_target(normalized),
        DatasetProvider::Zenodo => parse_zenodo_target(normalized),
    }
}

fn discover_hf(
    client: &Client,
    target: &DatasetTarget,
    token: Option<&str>,
    page_size: usize,
) -> Result<Vec<DatasetSummary>> {
    match target {
        DatasetTarget::Dataset(id) => Ok(vec![inspect_hf(client, id, token)?.summary]),
        DatasetTarget::Namespace(namespace) => list_hf_datasets(
            client,
            Some(("author", namespace)),
            "downloads",
            -1,
            token,
            page_size,
        ),
        DatasetTarget::Search(query) => list_hf_datasets(
            client,
            Some(("search", query)),
            "downloads",
            -1,
            token,
            page_size,
        ),
        DatasetTarget::Top => list_hf_datasets(client, None, "downloads", -1, token, page_size),
        DatasetTarget::Latest => {
            list_hf_datasets(client, None, "lastModified", -1, token, page_size)
        }
    }
}

fn list_hf_datasets(
    client: &Client,
    filter: Option<(&str, &str)>,
    sort: &str,
    direction: i8,
    token: Option<&str>,
    page_size: usize,
) -> Result<Vec<DatasetSummary>> {
    let mut page = 1usize;
    let mut datasets = Vec::new();

    loop {
        let mut url = format!(
            "{HF_BASE}/api/datasets?sort={}&direction={direction}&limit={}&page={page}",
            urlencoding::encode(sort),
            page_size
        );
        if let Some((key, value)) = filter {
            url.push('&');
            url.push_str(key);
            url.push('=');
            url.push_str(&urlencoding::encode(value));
        }

        let response = send_hf(client.get(&url), token)?;
        let status = response.status();
        if !status.is_success() {
            return Err(api_error(DatasetProvider::Hf, response));
        }

        let batch: Vec<HfDatasetResponse> = response.json()?;
        let count = batch.len();
        datasets.extend(batch.into_iter().map(hf_summary));
        if count < page_size {
            break;
        }
        page = page.saturating_add(1);
    }

    Ok(datasets)
}

fn inspect_hf(client: &Client, dataset_id: &str, token: Option<&str>) -> Result<DatasetDetails> {
    let url = format!("{HF_BASE}/api/datasets/{}", encode_repo_id(dataset_id));
    let response = send_hf(client.get(&url), token)?;
    let status = response.status();

    if status == StatusCode::NOT_FOUND {
        return Err(DataForgeError::NotFound {
            provider: DatasetProvider::Hf.slug(),
            target: dataset_id.to_string(),
        });
    }
    if !status.is_success() {
        return Err(api_error(DatasetProvider::Hf, response));
    }

    let data: HfDatasetResponse = response.json()?;
    let summary = hf_summary(data.clone());
    let files = data
        .siblings
        .into_iter()
        .filter(|file| !file.rfilename.trim().is_empty())
        .map(|file| DatasetFile {
            download_url: hf_file_url(&summary.id, &file.rfilename),
            path: file.rfilename,
            size: file.size,
            checksum: file.blob_id,
        })
        .collect();

    Ok(DatasetDetails { summary, files })
}

fn discover_zenodo(
    client: &Client,
    target: &DatasetTarget,
    page_size: usize,
) -> Result<Vec<DatasetSummary>> {
    match target {
        DatasetTarget::Dataset(id) => Ok(vec![inspect_zenodo(client, id)?.summary]),
        DatasetTarget::Namespace(value) | DatasetTarget::Search(value) => {
            list_zenodo(client, Some(value), "mostviewed", page_size)
        }
        DatasetTarget::Top => list_zenodo(client, None, "mostviewed", page_size),
        DatasetTarget::Latest => list_zenodo(client, None, "newest", page_size),
    }
}

fn list_zenodo(
    client: &Client,
    query: Option<&str>,
    sort: &str,
    page_size: usize,
) -> Result<Vec<DatasetSummary>> {
    let mut page = 1usize;
    let mut datasets = Vec::new();

    loop {
        let mut url = format!(
            "{ZENODO_API}/records?sort={}&size={}&page={page}",
            urlencoding::encode(sort),
            page_size
        );
        if let Some(query) = query {
            url.push_str("&q=");
            url.push_str(&urlencoding::encode(query));
        }

        let response = client.get(&url).send()?;
        let status = response.status();
        if !status.is_success() {
            return Err(api_error(DatasetProvider::Zenodo, response));
        }

        let batch: ZenodoSearchResponse = response.json()?;
        let count = batch.hits.hits.len();
        datasets.extend(batch.hits.hits.into_iter().map(zenodo_summary));
        if count < page_size {
            break;
        }
        page = page.saturating_add(1);
    }

    Ok(datasets)
}

fn inspect_zenodo(client: &Client, record_id: &str) -> Result<DatasetDetails> {
    let url = format!("{ZENODO_API}/records/{}", urlencoding::encode(record_id));
    let response = client.get(&url).send()?;
    let status = response.status();

    if status == StatusCode::NOT_FOUND {
        return Err(DataForgeError::NotFound {
            provider: DatasetProvider::Zenodo.slug(),
            target: record_id.to_string(),
        });
    }
    if !status.is_success() {
        return Err(api_error(DatasetProvider::Zenodo, response));
    }

    let data: ZenodoRecord = response.json()?;
    let summary = zenodo_summary(data.clone());
    let files = data
        .files
        .into_iter()
        .filter_map(|file| {
            let url = file.links.self_url.or(file.links.download)?;
            Some(DatasetFile {
                path: file.key,
                download_url: url,
                size: file.size,
                checksum: file.checksum,
            })
        })
        .collect();

    Ok(DatasetDetails { summary, files })
}

fn run_download_jobs(
    client: Client,
    token: Option<String>,
    concurrency: usize,
    jobs: VecDeque<DownloadJob>,
) -> Vec<std::result::Result<ArchivedFile, ArchiveFailure>> {
    let worker_count = concurrency.max(1).min(jobs.len().max(1));
    let jobs = Arc::new(Mutex::new(jobs));
    let results = Arc::new(Mutex::new(Vec::new()));
    let mut workers = Vec::with_capacity(worker_count);

    for _ in 0..worker_count {
        let jobs = Arc::clone(&jobs);
        let results = Arc::clone(&results);
        let client = client.clone();
        let token = token.clone();
        workers.push(thread::spawn(move || {
            loop {
                let job = {
                    let mut queue = jobs.lock().expect("download queue lock poisoned");
                    queue.pop_front()
                };
                let Some(job) = job else {
                    break;
                };
                let result = download_job(&client, token.as_deref(), job);
                results
                    .lock()
                    .expect("download results lock poisoned")
                    .push(result);
            }
        }));
    }

    for worker in workers {
        let _ = worker.join();
    }

    Arc::try_unwrap(results)
        .unwrap_or_else(|arc| {
            Mutex::new(arc.lock().expect("download results lock poisoned").clone())
        })
        .into_inner()
        .expect("download results lock poisoned")
}

fn download_job(
    client: &Client,
    token: Option<&str>,
    job: DownloadJob,
) -> std::result::Result<ArchivedFile, ArchiveFailure> {
    if let Some(parent) = job.local_path.parent() {
        if let Err(error) = fs::create_dir_all(parent) {
            return Err(job.failure(None, error));
        }
    }

    let mut request = client.get(&job.url);
    if let Some(token) = token {
        request = request.header("Authorization", format!("Bearer {token}"));
    }

    let mut response = match request.send() {
        Ok(response) => response,
        Err(error) => return Err(job.failure(None, error)),
    };
    let status = response.status();
    if !status.is_success() {
        let body = response
            .text()
            .unwrap_or_else(|_| "<unable to read body>".to_string());
        return Err(job.failure(None, format!("HTTP {status}: {body}")));
    }

    let tmp_path = partial_path(&job.local_path);
    let mut file = match File::create(&tmp_path) {
        Ok(file) => file,
        Err(error) => return Err(job.failure(None, error)),
    };
    let bytes_written = match io::copy(&mut response, &mut file) {
        Ok(bytes) => bytes,
        Err(error) => {
            let _ = fs::remove_file(&tmp_path);
            return Err(job.failure(None, error));
        }
    };
    if let Err(error) = fs::rename(&tmp_path, &job.local_path) {
        let _ = fs::remove_file(&tmp_path);
        return Err(job.failure(None, error));
    }

    Ok(ArchivedFile {
        dataset_id: job.dataset_id,
        source_path: job.source_path,
        local_path: job.local_path,
        bytes_written,
        skipped: false,
    })
}

#[derive(Debug, Clone)]
struct DownloadJob {
    dataset_id: String,
    source_path: String,
    url: String,
    local_path: PathBuf,
}

impl DownloadJob {
    fn failure(&self, file: Option<String>, error: impl fmt::Display) -> ArchiveFailure {
        ArchiveFailure {
            dataset_id: self.dataset_id.clone(),
            file: file.or_else(|| Some(self.source_path.clone())),
            message: error.to_string(),
        }
    }
}

#[derive(Debug, Clone, Deserialize)]
struct HfDatasetResponse {
    #[serde(default)]
    id: String,
    #[serde(default)]
    downloads: u64,
    #[serde(default)]
    likes: u64,
    #[serde(rename = "lastModified", default)]
    last_modified: Option<String>,
    #[serde(default)]
    tags: Vec<String>,
    #[serde(default)]
    siblings: Vec<HfSibling>,
    #[serde(default)]
    description: Option<String>,
}

#[derive(Debug, Clone, Deserialize)]
struct HfSibling {
    #[serde(default)]
    rfilename: String,
    #[serde(default)]
    size: Option<u64>,
    #[serde(rename = "blobId", default)]
    blob_id: Option<String>,
}

#[derive(Debug, Clone, Deserialize)]
struct ZenodoSearchResponse {
    hits: ZenodoHits,
}

#[derive(Debug, Clone, Deserialize)]
struct ZenodoHits {
    hits: Vec<ZenodoRecord>,
}

#[derive(Debug, Clone, Deserialize)]
struct ZenodoRecord {
    id: u64,
    #[serde(default)]
    metadata: ZenodoMetadata,
    #[serde(default)]
    updated: Option<String>,
    #[serde(default)]
    created: Option<String>,
    #[serde(default)]
    files: Vec<ZenodoFile>,
    #[serde(default)]
    stats: ZenodoStats,
    #[serde(default)]
    links: ZenodoRecordLinks,
}

#[derive(Debug, Clone, Default, Deserialize)]
struct ZenodoMetadata {
    #[serde(default)]
    title: Option<String>,
    #[serde(default)]
    description: Option<String>,
    #[serde(default)]
    keywords: Vec<String>,
}

#[derive(Debug, Clone, Default, Deserialize)]
struct ZenodoStats {
    #[serde(default)]
    downloads: Option<u64>,
}

#[derive(Debug, Clone, Default, Deserialize)]
struct ZenodoRecordLinks {
    #[serde(default)]
    html: Option<String>,
}

#[derive(Debug, Clone, Deserialize)]
struct ZenodoFile {
    key: String,
    #[serde(default)]
    size: Option<u64>,
    #[serde(default)]
    checksum: Option<String>,
    #[serde(default)]
    links: ZenodoFileLinks,
}

#[derive(Debug, Clone, Default, Deserialize)]
struct ZenodoFileLinks {
    #[serde(rename = "self", default)]
    self_url: Option<String>,
    #[serde(default)]
    download: Option<String>,
}

fn hf_summary(data: HfDatasetResponse) -> DatasetSummary {
    let file_count = if data.siblings.is_empty() {
        None
    } else {
        Some(data.siblings.len())
    };
    let size = data
        .siblings
        .iter()
        .try_fold(0u64, |acc, file| file.size.map(|size| acc + size));
    DatasetSummary {
        provider: DatasetProvider::Hf,
        id: data.id.clone(),
        title: Some(data.id.clone()),
        downloads: Some(data.downloads),
        likes: Some(data.likes),
        last_modified: data.last_modified,
        description: data.description,
        tags: data.tags,
        file_count,
        size,
        url: Some(format!("{HF_BASE}/datasets/{}", encode_repo_id(&data.id))),
    }
}

fn zenodo_summary(data: ZenodoRecord) -> DatasetSummary {
    let size = data
        .files
        .iter()
        .try_fold(0u64, |acc, file| file.size.map(|size| acc + size));
    let file_count = if data.files.is_empty() {
        None
    } else {
        Some(data.files.len())
    };
    DatasetSummary {
        provider: DatasetProvider::Zenodo,
        id: data.id.to_string(),
        title: data.metadata.title,
        downloads: data.stats.downloads,
        likes: None,
        last_modified: data.updated.or(data.created),
        description: data.metadata.description,
        tags: data.metadata.keywords,
        file_count,
        size,
        url: data.links.html,
    }
}

fn parse_hf_target(value: &str) -> Result<DatasetTarget> {
    if let Some(path) = strip_any_prefix_ignore_ascii_case(
        value,
        &[
            "https://huggingface.co/datasets/",
            "http://huggingface.co/datasets/",
            "https://www.huggingface.co/datasets/",
            "http://www.huggingface.co/datasets/",
            "https://hf.co/datasets/",
            "http://hf.co/datasets/",
            "https://www.hf.co/datasets/",
            "http://www.hf.co/datasets/",
        ],
    ) {
        return parse_dataset_path(path, true);
    }

    if value.contains(char::is_whitespace) {
        return Ok(DatasetTarget::Search(value.to_string()));
    }

    parse_dataset_path(value, false)
}

fn parse_zenodo_target(value: &str) -> Result<DatasetTarget> {
    if let Some(path) = strip_any_prefix_ignore_ascii_case(
        value,
        &[
            "https://zenodo.org/records/",
            "http://zenodo.org/records/",
            "https://www.zenodo.org/records/",
            "http://www.zenodo.org/records/",
            "https://zenodo.org/record/",
            "http://zenodo.org/record/",
            "https://www.zenodo.org/record/",
            "http://www.zenodo.org/record/",
        ],
    ) {
        let id = path.split(['?', '#', '/']).next().unwrap_or_default();
        return non_empty_target(id, DatasetTarget::Dataset);
    }

    if value.chars().all(|ch| ch.is_ascii_digit()) {
        return Ok(DatasetTarget::Dataset(value.to_string()));
    }
    Ok(DatasetTarget::Search(value.to_string()))
}

fn parse_dataset_path(path: &str, explicit_dataset_url: bool) -> Result<DatasetTarget> {
    let parts = path
        .split(['?', '#'])
        .next()
        .unwrap_or_default()
        .split('/')
        .filter(|part| !part.trim().is_empty())
        .collect::<Vec<_>>();
    match parts.len() {
        0 => Err(DataForgeError::InvalidTarget(path.to_string())),
        1 if explicit_dataset_url => Ok(DatasetTarget::Dataset(parts[0].to_string())),
        1 => Ok(DatasetTarget::Namespace(parts[0].to_string())),
        _ => Ok(DatasetTarget::Dataset(format!("{}/{}", parts[0], parts[1]))),
    }
}

fn non_empty_target(
    value: &str,
    build: impl FnOnce(String) -> DatasetTarget,
) -> Result<DatasetTarget> {
    let value = value.trim();
    if value.is_empty() {
        Err(DataForgeError::InvalidTarget(value.to_string()))
    } else {
        Ok(build(value.to_string()))
    }
}

fn send_hf(
    mut request: reqwest::blocking::RequestBuilder,
    token: Option<&str>,
) -> reqwest::Result<Response> {
    if let Some(token) = token {
        request = request.header("Authorization", format!("Bearer {token}"));
    }
    request.send()
}

fn api_error(provider: DatasetProvider, response: Response) -> DataForgeError {
    let status = response.status();
    let body = response
        .text()
        .unwrap_or_else(|_| "<unable to read body>".to_string());
    DataForgeError::ProviderApi {
        provider: provider.slug(),
        status: status.as_u16(),
        body,
    }
}

fn prepare_output(output: &Path) -> Result<()> {
    if output.exists() && !output.is_dir() {
        return Err(DataForgeError::OutputNotDirectory(output.to_path_buf()));
    }
    fs::create_dir_all(output)?;
    Ok(())
}

fn write_dataset_manifest(dataset_dir: &Path, details: &DatasetDetails) -> Result<()> {
    let manifest = serde_json::to_vec_pretty(details)?;
    fs::write(dataset_dir.join("dataforge-dataset.json"), manifest)?;
    Ok(())
}

fn dataset_matches_filter(dataset: &DatasetSummary, filter: Option<&str>) -> bool {
    let Some(filter) = filter.map(str::trim).filter(|filter| !filter.is_empty()) else {
        return true;
    };
    let filter = filter.to_lowercase();
    dataset.id.to_lowercase().contains(&filter)
        || dataset
            .title
            .as_deref()
            .is_some_and(|title| title.to_lowercase().contains(&filter))
}

fn dataset_dir_name(dataset_id: &str) -> String {
    sanitize_component(dataset_id).replace('/', "_")
}

fn safe_relative_path(path: &str) -> PathBuf {
    let mut clean = PathBuf::new();
    for part in path.split('/') {
        let part = sanitize_component(part);
        if part.is_empty() || part == "." || part == ".." {
            continue;
        }
        clean.push(part);
    }
    if clean.as_os_str().is_empty() {
        PathBuf::from("dataset.bin")
    } else {
        clean
    }
}

fn sanitize_component(value: &str) -> String {
    value
        .trim()
        .chars()
        .map(|ch| match ch {
            'a'..='z' | 'A'..='Z' | '0'..='9' | '-' | '_' | '.' | '/' => ch,
            _ => '_',
        })
        .collect::<String>()
        .trim_matches('_')
        .to_string()
}

fn partial_path(path: &Path) -> PathBuf {
    let mut partial = path.to_path_buf();
    let extension = path.extension().and_then(|ext| ext.to_str()).map_or_else(
        || "dataforge-part".to_string(),
        |ext| format!("{ext}.dataforge-part"),
    );
    partial.set_extension(extension);
    partial
}

fn hf_file_url(dataset_id: &str, path: &str) -> String {
    format!(
        "{HF_BASE}/datasets/{}/resolve/main/{}",
        encode_repo_id(dataset_id),
        encode_path(path)
    )
}

fn encode_repo_id(value: &str) -> String {
    encode_path(value)
}

fn encode_path(value: &str) -> String {
    value
        .split('/')
        .map(urlencoding::encode)
        .collect::<Vec<_>>()
        .join("/")
}

fn normalized_slug(value: &str) -> String {
    value
        .trim()
        .chars()
        .filter(|ch| !matches!(ch, '-' | '_' | ' ' | '\t' | '\n' | '\r'))
        .flat_map(char::to_lowercase)
        .collect()
}

fn strip_any_prefix_ignore_ascii_case<'a>(value: &'a str, prefixes: &[&str]) -> Option<&'a str> {
    prefixes.iter().find_map(|prefix| {
        value
            .get(..prefix.len())
            .is_some_and(|head| head.eq_ignore_ascii_case(prefix))
            .then(|| &value[prefix.len()..])
    })
}

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

    #[test]
    fn parses_hugging_face_dataset_urls() {
        assert_eq!(
            parse_dataset_target_for_provider(
                "https://huggingface.co/datasets/openai/gsm8k",
                DatasetProvider::Hf,
            )
            .unwrap(),
            DatasetTarget::Dataset("openai/gsm8k".to_string())
        );
        assert_eq!(
            parse_dataset_target_for_provider(
                "https://huggingface.co/datasets/squad",
                DatasetProvider::Hf,
            )
            .unwrap(),
            DatasetTarget::Dataset("squad".to_string())
        );
    }

    #[test]
    fn parses_hugging_face_collections_and_prefixes() {
        assert_eq!(
            parse_dataset_target_for_provider("openai", DatasetProvider::Hf).unwrap(),
            DatasetTarget::Namespace("openai".to_string())
        );
        assert_eq!(
            parse_dataset_target_for_provider("search:code data", DatasetProvider::Hf).unwrap(),
            DatasetTarget::Search("code data".to_string())
        );
        assert_eq!(
            parse_dataset_target_for_provider("dataset:squad", DatasetProvider::Hf).unwrap(),
            DatasetTarget::Dataset("squad".to_string())
        );
    }

    #[test]
    fn parses_zenodo_records() {
        assert_eq!(
            parse_dataset_target_for_provider(
                "https://zenodo.org/records/12345",
                DatasetProvider::Zenodo
            )
            .unwrap(),
            DatasetTarget::Dataset("12345".to_string())
        );
        assert_eq!(
            parse_dataset_target_for_provider("climate data", DatasetProvider::Zenodo).unwrap(),
            DatasetTarget::Search("climate data".to_string())
        );
    }

    #[test]
    fn parses_provider_aliases() {
        assert_eq!(
            DatasetProvider::from_slug("hugging-face"),
            Some(DatasetProvider::Hf)
        );
        assert_eq!(
            DatasetProvider::from_slug("zen"),
            Some(DatasetProvider::Zenodo)
        );
    }

    #[test]
    fn safe_paths_drop_traversal() {
        assert_eq!(safe_relative_path("../a/./b.csv"), PathBuf::from("a/b.csv"));
        assert_eq!(
            dataset_dir_name("owner/name with spaces"),
            "owner_name_with_spaces"
        );
    }
}