use std::io::{self, IsTerminal, Write};
use std::path::{Path, PathBuf};
use baseforge::{ArchiveOptions, BaseForge, DatasetSummary, DatasetTarget};
use crate::cli::{Cli, DatasetArchiveArgs};
use crate::constants::THESA_METADATA_DIR;
use crate::error::{Result, ThesaError};
use crate::manifest::{prepare_output_dir, write_archive_manifest};
use crate::selection::parse_indices;
use crate::terminal::TerminalSession;
use crate::types::DatasetProvider;
pub(crate) fn parse_dataset_target_for_provider(
input: &str,
provider: DatasetProvider,
) -> Result<DatasetTarget> {
baseforge::parse_dataset_target_for_provider(input, provider.into())
.map_err(|err| ThesaError::InvalidTarget(err.to_string()))
}
pub(crate) fn run_datasets_mode(
args: &mut Cli,
terminal_session: &mut Option<TerminalSession>,
) -> Result<()> {
*terminal_session = None;
let target_text = args.target.as_deref().unwrap_or_default().to_string();
let target = parse_dataset_target_for_provider(&target_text, args.dataset_provider)?;
if args.output == Path::new("./archives") {
args.output = PathBuf::from("./archives").join("datasets");
}
run_baseforge_archive(
&target_text,
&target,
args.dataset_provider,
&args.output,
args.concurrency,
args.skip_existing,
args.filter.clone(),
args.hf_token.clone(),
args.dry_run,
io::stdout().is_terminal(),
)
}
pub(crate) fn run_baseforge_archive_command(
args: &DatasetArchiveArgs,
dry_run: bool,
) -> Result<()> {
let target = parse_dataset_target_for_provider(&args.target, args.provider)?;
let hf_token = args
.hf_token
.clone()
.or_else(|| std::env::var("HF_TOKEN").ok());
run_baseforge_archive(
&args.target,
&target,
args.provider,
&args.archive_root,
args.concurrency,
args.skip_existing,
args.filter.clone(),
hf_token,
dry_run || args.dry_run,
io::stdout().is_terminal(),
)
}
#[allow(clippy::too_many_arguments)]
fn run_baseforge_archive(
target_text: &str,
target: &DatasetTarget,
provider: DatasetProvider,
output: &Path,
concurrency: usize,
skip_existing: bool,
filter: Option<String>,
hf_token: Option<String>,
dry_run: bool,
interactive: bool,
) -> Result<()> {
let forge = baseforge_client(provider, hf_token)?;
let discovered = forge
.discover(target)
.map_err(|err| baseforge_err("discover", err))?;
let discovered_count = discovered.len();
let mut selected = discovered
.into_iter()
.filter(|dataset| dataset_matches_filter(dataset, filter.as_deref()))
.collect::<Vec<_>>();
if dry_run {
print_dataset_plan(
provider,
output,
target_text,
concurrency,
skip_existing,
discovered_count,
&selected,
);
return Ok(());
}
if interactive && target.is_collection() && selected.len() > 1 {
selected = pick_datasets_interactively(&selected)?;
}
if selected.is_empty() {
println!("No datasets matched your request.");
return Ok(());
}
let output = prepare_output_dir(output)?;
let summary = forge
.archive_datasets(
&selected,
&ArchiveOptions {
output: output.clone(),
concurrency,
skip_existing,
filter: None,
},
)
.map_err(|err| baseforge_err("archive", err))?;
let manifest = write_archive_manifest(
&output,
&format!("dataset:{provider}"),
target_text,
summary.is_success(),
)?;
println!(
"\nCompleted: {} datasets archived, {} skipped, {} failed",
summary.archived,
summary.skipped,
summary.failed_dataset_count()
);
println!(
"Files: {} downloaded, {} bytes written",
summary.files_archived,
format_bytes(summary.bytes_written)
);
println!(
"Manifest: {}/manifest.json ({} files, {} bytes)",
THESA_METADATA_DIR, manifest.archive.file_count, manifest.archive.total_bytes
);
if !summary.failed.is_empty() {
eprintln!("Failures:");
for failure in summary.failed {
if let Some(file) = failure.file {
eprintln!(" - {}:{}: {}", failure.dataset_id, file, failure.message);
} else {
eprintln!(" - {}: {}", failure.dataset_id, failure.message);
}
}
return Err(ThesaError::Message(
"some dataset downloads failed".to_string(),
));
}
Ok(())
}
fn baseforge_client(provider: DatasetProvider, hf_token: Option<String>) -> Result<BaseForge> {
let mut forge = BaseForge::new(provider.into()).map_err(|err| baseforge_err("init", err))?;
if provider == DatasetProvider::Hf {
if let Some(token) = hf_token.filter(|token| !token.trim().is_empty()) {
forge = forge.with_token(token);
}
}
Ok(forge)
}
fn print_dataset_plan(
provider: DatasetProvider,
output: &Path,
target: &str,
concurrency: usize,
skip_existing: bool,
discovered_count: usize,
datasets: &[DatasetSummary],
) {
println!(
"Dry run: {} of {} datasets would be archived with BaseForge",
datasets.len(),
discovered_count
);
println!(
"Provider: {} ({})",
provider.label(),
provider.short_label()
);
println!("Target: {target}");
println!("Output: {}", output.display());
println!("Concurrency: {}", concurrency.max(1));
println!("Skip existing: {skip_existing}");
for dataset in datasets.iter().take(100) {
print_dataset_line(dataset);
}
if datasets.len() > 100 {
println!(" - ... {} more", datasets.len() - 100);
}
}
fn pick_datasets_interactively(datasets: &[DatasetSummary]) -> Result<Vec<DatasetSummary>> {
println!("Discovered {} datasets:", datasets.len());
for (index, dataset) in datasets.iter().enumerate().take(100) {
print!(" {:>3})", index + 1);
print_dataset_line(dataset);
}
if datasets.len() > 100 {
println!(" ... {} more", datasets.len() - 100);
}
println!("Select dataset indexes such as 1,3,5-8, or press Enter for all.");
print!("Selection: ");
io::stdout().flush()?;
let mut input = String::new();
io::stdin().read_line(&mut input).map_err(ThesaError::Io)?;
let input = input.trim();
if input.is_empty() || input.eq_ignore_ascii_case("all") {
return Ok(datasets.to_vec());
}
let indexes = parse_indices(input, datasets.len()).map_err(ThesaError::Message)?;
Ok(indexes
.into_iter()
.map(|index| datasets[index].clone())
.collect())
}
fn print_dataset_line(dataset: &DatasetSummary) {
let title = dataset.title.as_deref().unwrap_or(&dataset.id);
println!(
" - [{}] {} ({}, files: {}, size: {})",
dataset.provider.short_label(),
title,
format_downloads(dataset.downloads),
dataset
.file_count
.map_or_else(|| "?".to_string(), |count| count.to_string()),
dataset.size.map_or_else(|| "?".to_string(), format_bytes)
);
}
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 baseforge_err(context: &str, err: impl std::fmt::Display) -> ThesaError {
ThesaError::Message(format!("baseforge {context} failed: {err}"))
}
fn format_downloads(downloads: Option<u64>) -> String {
downloads.map_or_else(
|| "downloads: ?".to_string(),
|downloads| format!("downloads: {downloads}"),
)
}
fn format_bytes(bytes: u64) -> String {
const KIB: u64 = 1024;
const MIB: u64 = KIB * 1024;
const GIB: u64 = MIB * 1024;
if bytes >= GIB {
format!("{:.2} GiB", bytes as f64 / GIB as f64)
} else if bytes >= MIB {
format!("{:.2} MiB", bytes as f64 / MIB as f64)
} else if bytes >= KIB {
format!("{:.2} KiB", bytes as f64 / KIB as f64)
} else {
format!("{bytes} B")
}
}