use std::path::PathBuf; use clap::{Parser, Subcommand}; use tracing_subscriber::EnvFilter; use shanty_db::Database; use shanty_db::entities::download_queue::DownloadStatus; use shanty_db::queries; use shanty_dl::{ AudioFormat, BackendConfig, DownloadBackend, DownloadTarget, SearchSource, YtDlpBackend, download_single, run_queue, }; #[derive(Parser)] #[command(name = "shanty-dl", about = "Download music files for Shanty")] struct Cli { #[command(subcommand)] command: Commands, /// Database URL. Defaults to sqlite:///shanty/shanty.db?mode=rwc #[arg(long, global = true, env = "SHANTY_DATABASE_URL")] database: Option, /// Increase verbosity (-v info, -vv debug, -vvv trace). #[arg(short, long, global = true, action = clap::ArgAction::Count)] verbose: u8, } #[derive(Subcommand)] enum Commands { /// Download a single song by query or URL. Download { /// Search query or direct URL. query_or_url: String, /// Output audio format. #[arg(long, default_value = "opus")] format: String, /// Output directory for downloaded files. #[arg(long)] output: Option, /// Path to cookies.txt file for YouTube authentication. #[arg(long, env = "SHANTY_COOKIES")] cookies: Option, /// Search source (ytmusic or youtube). #[arg(long, default_value = "ytmusic")] search_source: String, /// Requests per hour limit. #[arg(long, default_value = "450")] rate_limit: u32, /// Preview what would be downloaded without doing it. #[arg(long)] dry_run: bool, }, /// Manage the download queue. Queue { #[command(subcommand)] action: QueueAction, }, } #[derive(Subcommand)] enum QueueAction { /// Process all pending items in the download queue. Process { /// Output audio format. #[arg(long, default_value = "opus")] format: String, /// Output directory for downloaded files. #[arg(long)] output: Option, /// Path to cookies.txt for YouTube authentication. #[arg(long, env = "SHANTY_COOKIES")] cookies: Option, /// Search source (ytmusic or youtube). #[arg(long, default_value = "ytmusic")] search_source: String, /// Requests per hour limit. #[arg(long, default_value = "450")] rate_limit: u32, /// Preview without downloading. #[arg(long)] dry_run: bool, }, /// Add an item to the download queue. Add { /// Search query for the song to download. query: String, }, /// List items in the download queue. List { /// Filter by status (pending, downloading, completed, failed, cancelled, all). #[arg(long, default_value = "all")] status: String, }, /// Retry all failed downloads. Retry, /// Sync wanted items from the watchlist to the download queue. Sync { /// Preview what would be enqueued without doing it. #[arg(long)] dry_run: bool, }, } fn default_database_url() -> String { let data_dir = dirs::data_dir() .unwrap_or_else(|| PathBuf::from(".")) .join("shanty"); std::fs::create_dir_all(&data_dir).ok(); let db_path = data_dir.join("shanty.db"); format!("sqlite://{}?mode=rwc", db_path.display()) } fn default_output_dir() -> PathBuf { let dir = dirs::data_dir() .unwrap_or_else(|| PathBuf::from(".")) .join("shanty") .join("downloads"); std::fs::create_dir_all(&dir).ok(); dir } fn make_backend( cookies: &Option, search_source: &str, rate_limit: u32, ) -> anyhow::Result { let source: SearchSource = search_source .parse() .map_err(|e: String| anyhow::anyhow!(e))?; // Bump rate limit if cookies are provided let effective_rate = if cookies.is_some() && rate_limit == 450 { tracing::info!("cookies provided — using authenticated rate limit (1800/hr)"); 1800 } else { rate_limit }; Ok(YtDlpBackend::new(effective_rate, source, cookies.clone())) } fn make_backend_config( format: &str, output: &Option, cookies: &Option, ) -> anyhow::Result { let fmt: AudioFormat = format.parse().map_err(|e: String| anyhow::anyhow!(e))?; Ok(BackendConfig { output_dir: output.clone().unwrap_or_else(default_output_dir), format: fmt, cookies_path: cookies.clone(), }) } #[tokio::main] async fn main() -> anyhow::Result<()> { let cli = Cli::parse(); // Set up tracing let filter = match cli.verbose { 0 => "warn", 1 => "info,shanty_dl=info", 2 => "info,shanty_dl=debug", _ => "debug,shanty_dl=trace", }; tracing_subscriber::fmt() .with_env_filter( EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(filter)), ) .init(); match cli.command { Commands::Download { query_or_url, format, output, cookies, search_source, rate_limit, dry_run, } => { let backend = make_backend(&cookies, &search_source, rate_limit)?; backend.check_available().await?; let config = make_backend_config(&format, &output, &cookies)?; // Determine if it's a URL or a search query let target = if query_or_url.starts_with("http://") || query_or_url.starts_with("https://") { DownloadTarget::Url(query_or_url) } else { DownloadTarget::Query(query_or_url) }; download_single(&backend, target, &config, dry_run).await?; } Commands::Queue { action } => { let database_url = cli.database.unwrap_or_else(default_database_url); let db = Database::new(&database_url).await?; match action { QueueAction::Process { format, output, cookies, search_source, rate_limit, dry_run, } => { let backend = make_backend(&cookies, &search_source, rate_limit)?; backend.check_available().await?; let config = make_backend_config(&format, &output, &cookies)?; if dry_run { println!("DRY RUN — no files will be downloaded"); } let stats = run_queue(db.conn(), &backend, &config, dry_run).await?; println!("\nQueue processing complete: {stats}"); } QueueAction::Add { query } => { let item = queries::downloads::enqueue(db.conn(), &query, None, "ytdlp").await?; println!("Added to queue: id={}, query=\"{}\"", item.id, item.query); } QueueAction::List { status } => { let filter = match status.to_lowercase().as_str() { "all" => None, "pending" => Some(DownloadStatus::Pending), "downloading" => Some(DownloadStatus::Downloading), "completed" => Some(DownloadStatus::Completed), "failed" => Some(DownloadStatus::Failed), "cancelled" => Some(DownloadStatus::Cancelled), _ => anyhow::bail!("unknown status: {status}"), }; let items = queries::downloads::list(db.conn(), filter).await?; if items.is_empty() { println!("Queue is empty."); } else { println!( "{:<5} {:<12} {:<6} {:<40} ERROR", "ID", "STATUS", "RETRY", "QUERY" ); for item in &items { println!( "{:<5} {:<12} {:<6} {:<40} {}", item.id, format!("{:?}", item.status), item.retry_count, truncate(&item.query, 40), item.error_message.as_deref().unwrap_or(""), ); } println!("\n{} items total", items.len()); } } QueueAction::Retry => { let failed = queries::downloads::list(db.conn(), Some(DownloadStatus::Failed)).await?; if failed.is_empty() { println!("No failed downloads to retry."); } else { for item in &failed { queries::downloads::retry_failed(db.conn(), item.id).await?; } println!("Requeued {} failed downloads.", failed.len()); } } QueueAction::Sync { dry_run } => { if dry_run { println!("DRY RUN — no items will be enqueued"); } let stats = shanty_dl::sync_wanted_to_queue(db.conn(), dry_run).await?; println!("\nSync complete: {stats}"); } } } } Ok(()) } fn truncate(s: &str, max: usize) -> String { if s.len() <= max { s.to_string() } else { format!("{}…", &s[..max - 1]) } }