use std::io::{self, IsTerminal, Write};
use std::path::{Path, PathBuf};
use aisling::{Loader, LoaderConfig, LoaderKind, LoaderProgress};
use baseforge::{
ArchiveOptions, ArchiveProgress, ArchiveProgressPhase, BaseForge, DatasetSummary, DatasetTarget,
};
use crossterm::{
cursor, execute,
terminal::{Clear, ClearType},
};
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)?;
println!(
"Archiving {} dataset(s) with BaseForge to {}",
selected.len(),
output.display()
);
println!("Concurrency: {}", concurrency.max(1));
io::stdout().flush()?;
let options = ArchiveOptions {
output: output.clone(),
concurrency,
skip_existing,
filter: None,
};
let mut progress = io::stderr()
.is_terminal()
.then(DatasetProgressRenderer::new);
let summary = forge
.archive_datasets_with_progress(&selected, &options, |snapshot| match progress.as_mut() {
Some(renderer) => renderer.render(&snapshot),
None => print_dataset_progress_line(&snapshot),
})
.map_err(|err| baseforge_err("archive", err))?;
if let Some(renderer) = progress.as_mut() {
renderer.finish();
}
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}"))
}
struct DatasetProgressRenderer {
stderr: io::Stderr,
loader: Loader,
use_color: bool,
width: usize,
tick: usize,
wrote: bool,
}
impl DatasetProgressRenderer {
fn new() -> Self {
let use_color = dataset_progress_color_enabled();
let width = dataset_terminal_width();
let config = LoaderConfig::default()
.with_width(dataset_loader_width(width))
.with_label("baseforge")
.with_unit("files")
.with_fraction(true);
let config = if use_color {
config
} else {
config.without_color()
};
let mut stderr = io::stderr();
let _ = execute!(stderr, cursor::Hide);
Self {
stderr,
loader: Loader::with_config(LoaderKind::Tqdm, config),
use_color,
width,
tick: 0,
wrote: false,
}
}
fn render(&mut self, progress: &ArchiveProgress) {
let current = progress.completed_files.min(progress.total_files.max(1)) as u64;
let total = progress.total_files.max(1) as u64;
let loader_line = if self.use_color {
self.loader
.frame(self.tick, LoaderProgress::from_counts(current, total))
.to_ansi_string()
} else {
self.loader
.line(self.tick, LoaderProgress::from_counts(current, total))
};
let line = dataset_progress_line(&loader_line, progress, self.width, self.use_color);
let _ = execute!(
self.stderr,
cursor::MoveToColumn(0),
Clear(ClearType::CurrentLine)
);
let _ = write!(
self.stderr,
"{}",
truncate_to_visible_width(&line, self.width.saturating_sub(1))
);
let _ = self.stderr.flush();
self.tick = self.tick.wrapping_add(1);
self.wrote = true;
}
fn finish(&mut self) {
if self.wrote {
let _ = writeln!(self.stderr);
}
let _ = execute!(self.stderr, cursor::Show);
let _ = self.stderr.flush();
}
}
impl Drop for DatasetProgressRenderer {
fn drop(&mut self) {
let _ = execute!(self.stderr, cursor::Show);
let _ = self.stderr.flush();
}
}
fn print_dataset_progress_line(progress: &ArchiveProgress) {
if matches!(progress.phase, ArchiveProgressPhase::Completed) {
return;
}
if progress.active_file.is_some() || progress.completed_files > 0 {
println!(
"Downloading {}/{} files (ok={}, failed={}, bytes={}, active={})",
progress.completed_files,
progress.total_files,
progress.succeeded_files,
progress.failed_files,
format_bytes(progress.bytes_written),
format_active_bytes(progress)
);
}
}
fn dataset_progress_line(
loader_line: &str,
progress: &ArchiveProgress,
width: usize,
use_color: bool,
) -> String {
let phase = colorize_progress(
progress.phase.as_str(),
dataset_phase_color(progress.phase),
use_color,
);
let stats = format!(
"ok={} fail={} total={}",
progress.succeeded_files,
progress.failed_files,
format_bytes(progress.bytes_written)
);
let active_stats = format_active_bytes(progress);
let fixed_width = visible_width(loader_line)
.saturating_add(progress.phase.as_str().len())
.saturating_add(stats.len())
.saturating_add(active_stats.len())
.saturating_add(26);
let active_budget = width.saturating_sub(fixed_width);
let active = progress
.active_file
.as_ref()
.filter(|_| active_budget >= 18)
.map(|file| {
let label = format!("{}:{}", file.dataset_id, file.source_path);
format!(" | {}", truncate_middle(&label, active_budget.min(72)))
})
.unwrap_or_default();
format!(
"{} | {} | active={} | {}{}",
loader_line.trim_end(),
phase,
active_stats,
stats,
active
)
}
fn format_active_bytes(progress: &ArchiveProgress) -> String {
match progress.active_total_bytes {
Some(total) if total > 0 => format!(
"{}/{}",
format_bytes(progress.active_bytes_written),
format_bytes(total)
),
_ => format_bytes(progress.active_bytes_written),
}
}
fn dataset_phase_color(phase: ArchiveProgressPhase) -> &'static str {
match phase {
ArchiveProgressPhase::Starting => "37;1",
ArchiveProgressPhase::Downloading => "36;1",
ArchiveProgressPhase::Completed => "32;1",
}
}
fn dataset_loader_width(terminal_width: usize) -> usize {
match terminal_width {
0..=79 => 28,
80..=119 => 34,
_ => 46,
}
}
fn dataset_terminal_width() -> usize {
std::env::var("COLUMNS")
.ok()
.and_then(|value| value.parse::<usize>().ok())
.filter(|value| *value >= 40)
.unwrap_or(120)
}
fn dataset_progress_color_enabled() -> bool {
std::env::var_os("NO_COLOR").is_none()
&& std::env::var("TERM")
.map(|term| term != "dumb")
.unwrap_or(true)
}
fn colorize_progress(text: &str, code: &str, enabled: bool) -> String {
if enabled {
format!("\x1b[{code}m{text}\x1b[0m")
} else {
text.to_string()
}
}
fn truncate_middle(input: &str, max_chars: usize) -> String {
let chars = input.chars().collect::<Vec<_>>();
if chars.len() <= max_chars {
return input.to_string();
}
let keep = max_chars.saturating_sub(3);
let left = keep / 2;
let right = keep.saturating_sub(left);
let mut output = chars.iter().take(left).collect::<String>();
output.push_str("...");
output.push_str(
&chars[chars.len().saturating_sub(right)..]
.iter()
.collect::<String>(),
);
output
}
fn truncate_to_visible_width(input: &str, max_width: usize) -> String {
if visible_width(input) <= max_width {
return input.to_string();
}
let mut output = String::new();
let mut visible = 0usize;
let mut chars = input.chars().peekable();
while let Some(ch) = chars.next() {
if ch == '\x1b' {
output.push(ch);
for next in chars.by_ref() {
output.push(next);
if next == 'm' {
break;
}
}
continue;
}
if visible + 1 >= max_width.saturating_sub(1) {
output.push_str("...");
break;
}
output.push(ch);
visible += 1;
}
if input.contains("\x1b[") {
output.push_str("\x1b[0m");
}
output
}
fn visible_width(input: &str) -> usize {
let mut width = 0usize;
let mut chars = input.chars().peekable();
while let Some(ch) = chars.next() {
if ch == '\x1b' {
for next in chars.by_ref() {
if next == 'm' {
break;
}
}
continue;
}
width += 1;
}
width
}
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")
}
}