use std::collections::BTreeSet;
use std::io::{self, IsTerminal, Write};
use std::path::{Path, PathBuf};
use std::sync::mpsc;
use std::thread;
use std::time::Duration;
use aisling::{Loader, LoaderConfig, LoaderKind, LoaderProgress};
use assetforge::{
AMBIENTCG_API_TERMS_URL, ArchiveOptions, ArchiveProgress, ArchiveProgressPhase, AssetBundle,
AssetForge, AssetSummary, AssetTarget, GAZEBO_FUEL_API_TERMS_URL,
GOOGLE_SCANNED_OBJECTS_LICENSE_URL, POLY_HAVEN_API_TERMS_URL, SKETCHFAB_API_TERMS_URL,
};
use crossterm::{
cursor,
event::{self, Event, KeyCode, KeyEventKind},
execute,
terminal::{Clear, ClearType},
};
use crate::cli::{AssetArchiveArgs, Cli};
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::{
RankedPickerEntry, TerminalSession, draw_full_frame, pick_ranked_results_tui, read_input_line,
};
use crate::types::{AssetFormat, AssetProvider, AssetResolution, Mode};
const GSO_ATTRIBUTION: &str = "Google Research via Gazebo Fuel";
pub(crate) fn parse_asset_target_for_provider(
input: &str,
provider: AssetProvider,
) -> Result<AssetTarget> {
assetforge::parse_asset_target_for_provider(input, provider.into())
.map_err(|error| ThesaError::InvalidTarget(error.to_string()))
}
pub(crate) fn run_assets_mode(
args: &mut Cli,
terminal_session: &mut Option<TerminalSession>,
) -> Result<()> {
let target_text = args.target.as_deref().unwrap_or_default().to_string();
let target = parse_asset_target_for_provider(&target_text, args.asset_provider)?;
if args.output == Path::new("./archives") {
args.output = PathBuf::from("./archives/assets");
}
let forge = assetforge_client(
args.asset_provider,
args.asset_limit,
args.asset_token.as_deref(),
)?;
let discovered = if let Some(session) = terminal_session.as_mut() {
discover_assets_tui(
&mut session.terminal,
forge.clone(),
target.clone(),
args.asset_provider,
&target_text,
)?
} else {
forge
.discover(&target)
.map_err(|error| assetforge_error("discover", error))?
};
let discovered_count = discovered.len();
let mut selected = filtered_assets(discovered, args.filter.as_deref());
let used_tui = terminal_session.is_some();
let mut selection_confirmed = false;
if target.is_collection() && selected.len() > 1 {
if let Some(session) = terminal_session.as_mut() {
selected = pick_assets_tui(&mut session.terminal, &selected)?;
selection_confirmed = true;
}
}
*terminal_session = None;
finish_assetforge_archive(
&forge,
&target_text,
&target,
args.asset_provider,
args.asset_format,
args.asset_resolution,
&args.output,
args.skip_existing,
args.asset_max_download_bytes,
args.dry_run,
!used_tui && io::stdout().is_terminal(),
discovered_count,
selected,
selection_confirmed,
)
}
pub(crate) fn run_assetforge_archive_command(args: &AssetArchiveArgs, dry_run: bool) -> Result<()> {
let target = parse_asset_target_for_provider(&args.target, args.provider)?;
run_assetforge_archive(
&args.target,
&target,
args.provider,
args.format,
args.resolution,
&args.archive_root,
args.skip_existing,
args.filter.as_deref(),
args.max_download_bytes,
args.limit,
args.token.as_deref(),
dry_run || args.dry_run,
io::stdout().is_terminal(),
)
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn run_assetforge_archive(
target_text: &str,
target: &AssetTarget,
provider: AssetProvider,
format: AssetFormat,
resolution: AssetResolution,
output: &Path,
skip_existing: bool,
filter: Option<&str>,
max_download_bytes: u64,
discovery_limit: usize,
api_token: Option<&str>,
dry_run: bool,
interactive: bool,
) -> Result<()> {
let forge = assetforge_client(provider, discovery_limit, api_token)?;
let discovered = forge
.discover(target)
.map_err(|error| assetforge_error("discover", error))?;
let discovered_count = discovered.len();
let selected = filtered_assets(discovered, filter);
finish_assetforge_archive(
&forge,
target_text,
target,
provider,
format,
resolution,
output,
skip_existing,
max_download_bytes,
dry_run,
interactive,
discovered_count,
selected,
false,
)
}
#[allow(clippy::too_many_arguments)]
fn finish_assetforge_archive(
forge: &AssetForge,
target_text: &str,
target: &AssetTarget,
provider: AssetProvider,
format: AssetFormat,
resolution: AssetResolution,
output: &Path,
skip_existing: bool,
max_download_bytes: u64,
dry_run: bool,
interactive: bool,
discovered_count: usize,
mut selected: Vec<AssetSummary>,
selection_confirmed: bool,
) -> Result<()> {
print_asset_provider_notice(provider);
if selected.is_empty() {
if selection_confirmed {
println!("No free 3D assets selected; nothing to archive.");
} else {
println!("No free 3D assets matched your request.");
}
return Ok(());
}
if dry_run {
let bundles = selected
.iter()
.map(|asset| {
let bundle = forge
.bundle(&asset.id, format.into(), resolution.into())
.map_err(|error| assetforge_error("resolve bundle", error))?;
enforce_bundle_limit(&asset.id, bundle.total_bytes, max_download_bytes)?;
Ok(bundle)
})
.collect::<Result<Vec<_>>>()?;
print_asset_plan(
provider,
format,
resolution,
output,
target_text,
skip_existing,
max_download_bytes,
discovered_count,
&selected,
&bundles,
);
return Ok(());
}
if interactive && target.is_collection() && selected.len() > 1 {
selected = pick_assets_interactively(&selected)?;
}
let output = prepare_output_dir(output)?;
println!(
"Archiving {} free 3D asset(s) with AssetForge to {}",
selected.len(),
output.display()
);
println!(
"Source: {} | License: {} | Format: {format} | Resolution: {resolution}",
provider.label(),
asset_license(provider),
);
match provider {
AssetProvider::Polyhaven => {
println!("Poly Haven API terms: {POLY_HAVEN_API_TERMS_URL}")
}
AssetProvider::Ambientcg => {
println!("ambientCG API terms: {AMBIENTCG_API_TERMS_URL}")
}
AssetProvider::Sketchfab => {
println!("Sketchfab API terms: {SKETCHFAB_API_TERMS_URL}")
}
AssetProvider::GoogleScannedObjects => {
println!("Attribution: {GSO_ATTRIBUTION}");
println!("Google Scanned Objects license: {GOOGLE_SCANNED_OBJECTS_LICENSE_URL}");
println!("Gazebo Fuel API documentation: {GAZEBO_FUEL_API_TERMS_URL}");
}
}
io::stdout().flush()?;
let options = ArchiveOptions {
output: output.clone(),
format: format.into(),
resolution: resolution.into(),
skip_existing,
max_download_bytes,
};
let mut renderer = io::stderr().is_terminal().then(AssetProgressRenderer::new);
let mut last_plain_completed = 0usize;
let mut last_plain_bytes = 0u64;
let summary = forge
.archive_assets_with_progress(&selected, &options, |snapshot| {
if let Some(renderer) = renderer.as_mut() {
renderer.render(&snapshot);
} else {
let completed_changed = snapshot.completed_assets != last_plain_completed;
let bytes_advanced = snapshot.bytes_written < last_plain_bytes
|| snapshot.bytes_written.saturating_sub(last_plain_bytes) >= 8 * 1024 * 1024;
if completed_changed || bytes_advanced {
print_plain_progress(&snapshot);
last_plain_completed = snapshot.completed_assets;
last_plain_bytes = snapshot.bytes_written;
}
}
})
.map_err(|error| assetforge_error("archive", error))?;
if let Some(renderer) = renderer.as_mut() {
renderer.finish();
}
let failed_ids = summary
.failed
.iter()
.map(|failure| failure.asset_id.as_str())
.collect::<BTreeSet<_>>();
let mut manifest_count = 0usize;
let mut manifest_file_count = 0usize;
let mut manifest_total_bytes = 0u64;
for asset in &selected {
if failed_ids.contains(asset.id.as_str()) {
continue;
}
let asset_output = output.join(&asset.id);
if !asset_output.is_dir() {
continue;
}
let manifest =
write_archive_manifest(&asset_output, &format!("asset:{provider}"), &asset.id, true)?;
manifest_count += 1;
manifest_file_count += manifest.archive.file_count;
manifest_total_bytes = manifest_total_bytes.saturating_add(manifest.archive.total_bytes);
}
println!(
"\nCompleted: {} assets archived, {} skipped, {} failed",
summary.archived,
summary.skipped,
summary.failed.len()
);
println!(
"Files: {} downloaded, {} written",
summary.files_archived,
format_bytes(summary.bytes_written)
);
println!(
"Manifests: {manifest_count} asset-local {THESA_METADATA_DIR}/manifest.json sidecars ({manifest_file_count} files, {manifest_total_bytes} bytes)"
);
if !summary.failed.is_empty() {
eprintln!("Failures:");
for failure in summary.failed {
eprintln!(" - {}: {}", failure.asset_id, failure.message);
}
return Err(ThesaError::Message(
"some free 3D asset downloads failed".to_string(),
));
}
Ok(())
}
fn assetforge_client(
provider: AssetProvider,
discovery_limit: usize,
api_token: Option<&str>,
) -> Result<AssetForge> {
let api_token = validated_asset_token(provider, api_token)?;
let mut forge = AssetForge::with_user_agent(
provider.into(),
concat!(
"thesa/",
env!("CARGO_PKG_VERSION"),
" (AssetForge; https://github.com/Tknott95/Thesa)"
),
)
.map_err(|error| assetforge_error("init", error))?
.with_discovery_limit(discovery_limit);
if let Some(token) = api_token {
forge = forge.with_api_token(token);
}
Ok(forge)
}
fn print_asset_provider_notice(provider: AssetProvider) {
match provider {
AssetProvider::Polyhaven => eprintln!(
"Poly Haven API notice: reasonable non-commercial use unless separately licensed; {POLY_HAVEN_API_TERMS_URL}"
),
AssetProvider::Ambientcg => {
eprintln!("ambientCG API notice: CC0 assets; {AMBIENTCG_API_TERMS_URL}")
}
AssetProvider::Sketchfab => eprintln!(
"Sketchfab notice: authenticated, user-directed CC0 downloads with creator attribution; {SKETCHFAB_API_TERMS_URL}"
),
AssetProvider::GoogleScannedObjects => eprintln!(
"Google Scanned Objects notice: CC BY 4.0 attribution required ({GSO_ATTRIBUTION}); {GOOGLE_SCANNED_OBJECTS_LICENSE_URL}; API documentation: {GAZEBO_FUEL_API_TERMS_URL}"
),
}
}
fn discover_assets_tui(
terminal: &mut scrin::Terminal,
forge: AssetForge,
target: AssetTarget,
provider: AssetProvider,
target_text: &str,
) -> Result<Vec<AssetSummary>> {
let loading_label = format!(
"discovering {target_text} with {} through AssetForge",
provider.label()
);
let (tx, rx) = mpsc::channel();
thread::spawn(move || {
let result = forge.discover(&target);
let _ = tx.send(result);
});
let mut tick = 0usize;
loop {
if let Ok(result) = rx.try_recv() {
return result.map_err(|error| assetforge_error("discover", error));
}
draw_full_frame(terminal, |frame| {
crate::draw_loading_tui(frame, &loading_label, tick, Mode::Assets, None)
})?;
tick = tick.wrapping_add(1);
if event::poll(Duration::from_millis(80))? {
let Event::Key(key) = event::read()? else {
continue;
};
if key.kind == KeyEventKind::Press && key.code == KeyCode::Esc {
return Err(ThesaError::Message("aborted by user".to_string()));
}
}
}
}
fn filtered_assets(assets: Vec<AssetSummary>, filter: Option<&str>) -> Vec<AssetSummary> {
assets
.into_iter()
.filter(|asset| asset_matches_filter(asset, filter))
.collect()
}
fn pick_assets_tui(
terminal: &mut scrin::Terminal,
assets: &[AssetSummary],
) -> Result<Vec<AssetSummary>> {
let entries = assets
.iter()
.map(|asset| RankedPickerEntry {
provider: asset.provider.short_label().to_string(),
title: asset.name.clone(),
license: asset.rights.license.clone(),
metrics: format!(
"{} downloads | {} polygons",
crate::format_downloads(asset.download_count),
crate::format_downloads(asset.polycount)
),
details: vec![
("id".to_string(), asset.id.clone()),
("provider".to_string(), asset.provider.label().to_string()),
("attribution".to_string(), asset.rights.attribution.clone()),
],
})
.collect::<Vec<_>>();
let indexes = pick_ranked_results_tui(terminal, "asset picker", "assets", &entries)?;
Ok(indexes
.into_iter()
.map(|index| assets[index].clone())
.collect())
}
fn validated_asset_token(provider: AssetProvider, api_token: Option<&str>) -> Result<Option<&str>> {
if provider != AssetProvider::Sketchfab {
return Ok(None);
}
api_token
.map(str::trim)
.filter(|token| !token.is_empty())
.map(Some)
.ok_or_else(|| {
ThesaError::Message(
"Sketchfab archive cancelled: a nonblank API token is required".to_string(),
)
})
}
#[allow(clippy::too_many_arguments)]
fn print_asset_plan(
provider: AssetProvider,
format: AssetFormat,
resolution: AssetResolution,
output: &Path,
target: &str,
skip_existing: bool,
max_download_bytes: u64,
discovered_count: usize,
assets: &[AssetSummary],
bundles: &[AssetBundle],
) {
println!(
"Dry run: {} of {} free 3D assets selected as archive candidates with AssetForge",
assets.len(),
discovered_count
);
println!(
"Provider: {} ({})",
provider.label(),
provider.short_label()
);
println!(
"Attribution: {} | License: {}",
provider.label(),
asset_license(provider)
);
match provider {
AssetProvider::Polyhaven => println!("API terms: {POLY_HAVEN_API_TERMS_URL}"),
AssetProvider::Ambientcg => println!("API terms: {AMBIENTCG_API_TERMS_URL}"),
AssetProvider::Sketchfab => println!("API terms: {SKETCHFAB_API_TERMS_URL}"),
AssetProvider::GoogleScannedObjects => {
println!("Attribution: {GSO_ATTRIBUTION}");
println!("License: {GOOGLE_SCANNED_OBJECTS_LICENSE_URL}");
println!("API documentation: {GAZEBO_FUEL_API_TERMS_URL}");
}
}
println!("Target: {target}");
println!("Bundle: {resolution} {format}");
println!("Output: {}", output.display());
if skip_existing {
println!(
"Skip existing: true (provider manifest and checksum verification occurs during archive; candidates may be skipped)"
);
} else {
println!("Skip existing: false");
}
println!("Maximum bundle bytes: {max_download_bytes}");
for (asset, bundle) in assets.iter().zip(bundles) {
print_asset_line(asset, Some(bundle));
}
}
fn pick_assets_interactively(assets: &[AssetSummary]) -> Result<Vec<AssetSummary>> {
println!("Discovered {} ranked free 3D assets:", assets.len());
for (index, asset) in assets.iter().enumerate() {
print!(" {:>3})", index + 1);
print_asset_line(asset, None);
}
println!("Select asset indexes such as 1,3,5-8, or press Enter for all.");
print!("Selection: ");
io::stdout().flush()?;
let input = read_input_line(&mut io::stdin().lock())?;
let input = input.trim();
if input.is_empty() || input.eq_ignore_ascii_case("all") {
return Ok(assets.to_vec());
}
let indexes = parse_indices(input, assets.len()).map_err(ThesaError::Message)?;
Ok(indexes
.into_iter()
.map(|index| assets[index].clone())
.collect())
}
fn print_asset_line(asset: &AssetSummary, bundle: Option<&AssetBundle>) {
let bundle = bundle
.map(|bundle| {
format!(
", {} {} {} bundle",
format_bytes(bundle.total_bytes),
bundle.resolution,
bundle.format
)
})
.unwrap_or_default();
println!(
" - [{} / {}] {} ({}, {} downloads, {} polygons{bundle}) | {}",
asset.rights.attribution,
asset.rights.license,
asset.name,
asset.id,
asset.download_count,
asset.polycount,
asset.source_url
);
}
fn asset_matches_filter(asset: &AssetSummary, filter: Option<&str>) -> bool {
let Some(filter) = filter.map(str::trim).filter(|filter| !filter.is_empty()) else {
return true;
};
let filter = filter.to_ascii_lowercase();
asset.id.to_ascii_lowercase().contains(&filter)
|| asset.name.to_ascii_lowercase().contains(&filter)
|| asset.description.to_ascii_lowercase().contains(&filter)
|| asset
.categories
.iter()
.chain(&asset.tags)
.any(|value| value.to_ascii_lowercase().contains(&filter))
}
fn assetforge_error(context: &str, error: impl std::fmt::Display) -> ThesaError {
ThesaError::Message(format!("assetforge {context} failed: {error}"))
}
fn asset_license(provider: AssetProvider) -> &'static str {
match provider {
AssetProvider::GoogleScannedObjects => "CC-BY-4.0",
AssetProvider::Polyhaven | AssetProvider::Ambientcg | AssetProvider::Sketchfab => "CC0-1.0",
}
}
fn enforce_bundle_limit(asset_id: &str, bundle_bytes: u64, max_download_bytes: u64) -> Result<()> {
if max_download_bytes > 0 && bundle_bytes > max_download_bytes {
Err(ThesaError::Message(format!(
"assetforge dry run rejected '{asset_id}': bundle is {bundle_bytes} bytes, exceeding the {max_download_bytes} byte limit"
)))
} else {
Ok(())
}
}
struct AssetProgressRenderer {
stderr: io::Stderr,
loader: Loader,
tick: usize,
wrote: bool,
}
impl AssetProgressRenderer {
fn new() -> Self {
let mut stderr = io::stderr();
let _ = execute!(stderr, cursor::Hide);
Self {
stderr,
loader: Loader::with_config(
LoaderKind::Tqdm,
LoaderConfig::default()
.with_width(28)
.with_label("assetforge")
.with_unit("assets")
.with_fraction(true),
),
tick: 0,
wrote: false,
}
}
fn render(&mut self, progress: &ArchiveProgress) {
let total = progress.total_assets.max(1) as u64;
let current = progress.completed_assets.min(progress.total_assets) as u64;
let loader = self
.loader
.frame(self.tick, LoaderProgress::from_counts(current, total))
.to_ansi_string();
let active_asset = progress.active_asset.as_deref().unwrap_or("assetforge");
let active = progress.active_file.as_deref().unwrap_or("finalizing");
let line = format!(
"{loader} | {} | {} / {}",
format_active_bytes(progress),
truncate(active_asset, 28),
truncate(active, 48)
);
let _ = execute!(
self.stderr,
cursor::MoveToColumn(0),
Clear(ClearType::CurrentLine)
);
let _ = write!(self.stderr, "{line}");
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 AssetProgressRenderer {
fn drop(&mut self) {
let _ = execute!(self.stderr, cursor::Show);
}
}
fn print_plain_progress(progress: &ArchiveProgress) {
if progress.phase != ArchiveProgressPhase::Starting {
if let (Some(asset), Some(file)) = (&progress.active_asset, &progress.active_file) {
println!(
"AssetForge downloading {asset}/{}: {}",
truncate(file, 64),
format_active_bytes(progress)
);
return;
}
let skipped = progress
.completed_assets
.saturating_sub(progress.succeeded_assets + progress.failed_assets);
println!(
"AssetForge {}/{} assets (ok={}, skipped={}, failed={}, bytes={})",
progress.completed_assets,
progress.total_assets,
progress.succeeded_assets,
skipped,
progress.failed_assets,
format_bytes(progress.bytes_written)
);
}
}
fn format_active_bytes(progress: &ArchiveProgress) -> String {
match progress.active_total_bytes {
Some(total) => format!(
"{}/{}",
format_bytes(progress.active_bytes_written),
format_bytes(total)
),
None => format_bytes(progress.bytes_written),
}
}
fn format_bytes(bytes: u64) -> String {
const UNITS: [&str; 5] = ["B", "KiB", "MiB", "GiB", "TiB"];
let mut value = bytes as f64;
let mut unit = 0usize;
while value >= 1024.0 && unit + 1 < UNITS.len() {
value /= 1024.0;
unit += 1;
}
if unit == 0 {
format!("{bytes} B")
} else {
format!("{value:.2} {}", UNITS[unit])
}
}
fn truncate(value: &str, max_chars: usize) -> String {
let count = value.chars().count();
if count <= max_chars {
value.to_string()
} else {
let keep = max_chars.saturating_sub(3);
format!("{}...", value.chars().take(keep).collect::<String>())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parses_assetforge_targets() {
assert_eq!(
parse_asset_target_for_provider(
"https://polyhaven.com/a/wooden_stool_02",
AssetProvider::Polyhaven
)
.unwrap(),
AssetTarget::Asset("wooden_stool_02".to_string())
);
assert!(parse_asset_target_for_provider("", AssetProvider::Polyhaven).is_err());
assert_eq!(
parse_asset_target_for_provider("search:chair", AssetProvider::GoogleScannedObjects)
.unwrap(),
AssetTarget::Search("chair".to_string())
);
assert_eq!(
assetforge::AssetProvider::from(AssetProvider::GoogleScannedObjects),
assetforge::AssetProvider::GoogleScannedObjects
);
}
#[test]
fn truncates_asset_progress_paths() {
assert_eq!(truncate("short", 10), "short");
assert_eq!(truncate("abcdefghijkl", 8), "abcde...");
}
#[test]
fn dry_run_bundle_limit_rejects_false_plan() {
assert!(enforce_bundle_limit("chair", 2, 1).is_err());
assert!(enforce_bundle_limit("chair", 2, 0).is_ok());
assert!(enforce_bundle_limit("chair", 2, 2).is_ok());
}
#[test]
fn sketchfab_requires_a_nonblank_token_before_provider_work() {
assert!(validated_asset_token(AssetProvider::Sketchfab, None).is_err());
assert!(validated_asset_token(AssetProvider::Sketchfab, Some(" ")).is_err());
assert_eq!(
validated_asset_token(AssetProvider::Sketchfab, Some(" token ")).unwrap(),
Some("token")
);
assert_eq!(
validated_asset_token(AssetProvider::Polyhaven, Some("ignored")).unwrap(),
None
);
}
}