use std::collections::BTreeSet;
use std::io::{self, IsTerminal, Write};
use std::path::{Path, PathBuf};
use aisling::{Loader, LoaderConfig, LoaderKind, LoaderProgress};
use assetforge::{
ArchiveOptions, ArchiveProgress, ArchiveProgressPhase, AssetForge, AssetSummary, AssetTarget,
POLY_HAVEN_API_TERMS_URL,
};
use crossterm::{
cursor, 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::TerminalSession;
use crate::types::{AssetFormat, AssetProvider, AssetResolution};
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<()> {
*terminal_session = None;
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");
}
run_assetforge_archive(
&target_text,
&target,
args.asset_provider,
args.asset_format,
args.asset_resolution,
&args.output,
args.skip_existing,
args.filter.as_deref(),
args.asset_max_download_bytes,
args.dry_run,
io::stdout().is_terminal(),
)
}
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,
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,
dry_run: bool,
interactive: bool,
) -> Result<()> {
let 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))?;
eprintln!(
"Poly Haven API notice: reasonable non-commercial use unless separately licensed; {POLY_HAVEN_API_TERMS_URL}"
);
let discovered = forge
.discover(target)
.map_err(|error| assetforge_error("discover", error))?;
let discovered_count = discovered.len();
let mut selected = discovered
.into_iter()
.filter(|asset| asset_matches_filter(asset, filter))
.collect::<Vec<_>>();
if dry_run {
let bundle_sizes = 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.total_bytes)
})
.collect::<Result<Vec<_>>>()?;
print_asset_plan(
provider,
format,
resolution,
output,
target_text,
skip_existing,
max_download_bytes,
discovered_count,
&selected,
&bundle_sizes,
);
return Ok(());
}
if interactive && target.is_collection() && selected.len() > 1 {
selected = pick_assets_interactively(&selected)?;
}
if selected.is_empty() {
println!("No free 3D assets matched your request.");
return Ok(());
}
let output = prepare_output_dir(output)?;
println!(
"Archiving {} free 3D asset(s) with AssetForge to {}",
selected.len(),
output.display()
);
println!(
"Source: {} | License: CC0-1.0 | Format: {format} | Resolution: {resolution}",
provider.label()
);
println!("Poly Haven API terms: {POLY_HAVEN_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 summary = forge
.archive_assets_with_progress(&selected, &options, |snapshot| {
if let Some(renderer) = renderer.as_mut() {
renderer.render(&snapshot);
} else if snapshot.completed_assets != last_plain_completed {
print_plain_progress(&snapshot);
last_plain_completed = snapshot.completed_assets;
}
})
.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(())
}
#[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],
bundle_sizes: &[u64],
) {
println!(
"Dry run: {} of {} free 3D assets would be archived with AssetForge",
assets.len(),
discovered_count
);
println!(
"Provider: {} ({})",
provider.label(),
provider.short_label()
);
println!("Attribution: Poly Haven | License: CC0-1.0");
println!("API terms: {POLY_HAVEN_API_TERMS_URL}");
println!("Target: {target}");
println!("Bundle: {resolution} {format}");
println!("Output: {}", output.display());
println!("Skip existing: {skip_existing}");
println!("Maximum bundle bytes: {max_download_bytes}");
for (asset, bundle_size) in assets.iter().zip(bundle_sizes) {
print_asset_line(asset, Some(*bundle_size));
}
}
fn pick_assets_interactively(assets: &[AssetSummary]) -> Result<Vec<AssetSummary>> {
println!("Discovered {} Poly Haven assets (CC0):", 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 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(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_size: Option<u64>) {
let bundle = bundle_size
.map(|size| format!(", {} bundle", format_bytes(size)))
.unwrap_or_default();
println!(
" - [{} / {}] {} ({}, {} polygons{bundle}) | {}",
asset.rights.attribution,
asset.rights.license,
asset.name,
asset.id,
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 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 = progress.active_file.as_deref().unwrap_or("finalizing");
let line = format!(
"{loader} | {} | {}",
format_active_bytes(progress),
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 {
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());
}
#[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());
}
}