mod airsonic; mod db; mod filesystem; mod lastfm; mod metadata; mod mpd; mod playlist; mod tui; use std::env; use std::path::{Path, PathBuf}; use clap::{Parser, Subcommand}; use rand::prelude::*; use playlist::Candidate; #[derive(Parser)] #[command(name = "drift")] struct Cli { #[command(subcommand)] command: Command, } #[derive(Subcommand)] enum Command { /// Index a music directory Index { /// Print progress #[arg(short)] verbose: bool, /// Re-index already indexed artists #[arg(short)] force: bool, /// Music directory to index directory: String, }, /// Build a playlist from similar artists Build { /// Print track scores #[arg(short)] verbose: bool, /// Queue in MPD #[arg(short, conflicts_with = "airsonic")] mpd: bool, /// Create Airsonic playlist #[arg(short, conflicts_with = "mpd")] airsonic: bool, /// Interleave artists evenly #[arg(short, conflicts_with = "random")] shuffle: bool, /// Fully randomize order #[arg(short, conflicts_with = "shuffle")] random: bool, /// Number of tracks #[arg(short = 'n', default_value_t = 20, value_parser = parse_positive_usize)] count: usize, /// Popularity bias (0=no preference, 10=heavy popular bias) #[arg(short, default_value_t = 5, value_parser = clap::value_parser!(u8).range(0..=10))] popularity: u8, /// Artist name to seed (or pick interactively) artist: Option, }, } struct BuildOptions { verbose: bool, mpd: bool, airsonic: bool, shuffle: bool, random: bool, count: usize, popularity_bias: u8, } fn parse_positive_usize(s: &str) -> Result { let n: usize = s.parse().map_err(|e| format!("{e}"))?; if n == 0 { return Err("value must be greater than 0".to_string()); } Ok(n) } fn db_path() -> PathBuf { let data_dir = env::var("XDG_DATA_HOME") .map(PathBuf::from) .unwrap_or_else(|_| { let home = env::var("HOME").expect("HOME not set"); PathBuf::from(home).join(".local/share") }); let dir = data_dir.join("drift"); std::fs::create_dir_all(&dir).expect("failed to create data directory"); dir.join("drift.db") } fn main() { let cli = Cli::parse(); match cli.command { Command::Index { verbose, force, directory } => { cmd_index(verbose, force, &directory); } Command::Build { verbose, mpd, airsonic, shuffle, random, count, popularity, artist } => { let opts = BuildOptions { verbose, mpd, airsonic, shuffle, random, count, popularity_bias: popularity, }; cmd_build(opts, artist.as_deref()); } } } fn cmd_index(verbose: bool, force: bool, directory: &str) { dotenvy::dotenv().ok(); let api_key = env::var("LASTFM_API_KEY").unwrap_or_default(); if api_key.is_empty() { eprintln!("Error: LASTFM_API_KEY not set"); std::process::exit(1); } let conn = db::open(&db_path()).expect("failed to open database"); let lastfm = lastfm::LastfmClient::new(api_key); let dir = Path::new(directory); for path in filesystem::walk_music_files(dir) { let tags = match metadata::read_tags(&path, &[ metadata::Tag::ArtistMbid, metadata::Tag::TrackMbid, metadata::Tag::ArtistName, metadata::Tag::TrackTitle, ]) { Ok(t) => t, Err(e) => { eprintln!("{}: could not read tags: {e}", path.display()); continue; } }; let Some(artist_mbid) = tags[0].clone() else { continue }; let recording_mbid = tags[1].clone(); let artist_name = tags[2].clone(); let track_title = tags[3].clone(); let already_indexed = match db::artist_exists(&conn, &artist_mbid) { Ok(exists) => exists, Err(e) => { eprintln!("DB error checking artist {artist_mbid}: {e}"); continue; } }; let display_name = artist_name.as_deref().unwrap_or(&artist_mbid); if !already_indexed || force { if verbose { println!("Indexing {display_name}..."); } match lastfm.get_similar_artists(&artist_mbid, artist_name.as_deref()) { Ok(similar) => { if let Err(e) = db::insert_artist_with_similar( &conn, &artist_mbid, artist_name.as_deref(), &similar, ) { eprintln!("DB error inserting artist {artist_mbid}: {e}"); continue; } } Err(e) => { eprintln!("Last.fm error for {artist_mbid}: {e}"); continue; } } match lastfm.get_top_tracks(&artist_mbid, artist_name.as_deref()) { Ok(top_tracks) => { if let Err(e) = db::insert_top_tracks(&conn, &artist_mbid, &top_tracks) { eprintln!("DB error inserting top tracks for {display_name}: {e}"); } } Err(e) => { eprintln!("Last.fm top tracks error for {display_name}: {e}"); } } } else if verbose { println!("Skipping {display_name} (already indexed)"); } let path_str = path.to_string_lossy(); if let Err(e) = db::insert_track(&conn, &path_str, &artist_mbid, recording_mbid.as_deref(), track_title.as_deref()) { eprintln!("DB error inserting track {}: {e}", path.display()); } } } fn resolve_artist(artists: &[(String, String)], query: &str) -> Option<(String, String)> { let q = query.to_lowercase(); // Tier 1: exact case-insensitive match for (mbid, name) in artists { if name.to_lowercase() == q { return Some((mbid.clone(), name.clone())); } } // Tier 2: contains case-insensitive for (mbid, name) in artists { if name.to_lowercase().contains(&q) { return Some((mbid.clone(), name.clone())); } } // Tier 3: subsequence fuzzy match for (mbid, name) in artists { if tui::fuzzy_match(&q, name) { return Some((mbid.clone(), name.clone())); } } None } fn cmd_build(opts: BuildOptions, artist: Option<&str>) { dotenvy::dotenv().ok(); let conn = db::open(&db_path()).expect("failed to open database"); let artists = match db::get_all_artists(&conn) { Ok(a) => a, Err(e) => { eprintln!("DB error: {e}"); std::process::exit(1); } }; if artists.is_empty() { eprintln!("No artists in database. Run 'index' first."); std::process::exit(1); } let (artist_mbid, seed_name) = if let Some(query) = artist { match resolve_artist(&artists, query) { Some((mbid, name)) => { eprintln!("Matched: {name}"); (mbid, name) } None => { eprintln!("No artist matching \"{query}\""); std::process::exit(1); } } } else { match tui::run_artist_picker(&artists) { Some(selection) => selection, None => std::process::exit(0), } }; build_playlist(&conn, &artist_mbid, &seed_name, &opts); } fn build_playlist( conn: &rusqlite::Connection, artist_mbid: &str, seed_name: &str, opts: &BuildOptions, ) { let similar = match db::get_available_similar_artists(conn, artist_mbid) { Ok(a) => a, Err(e) => { eprintln!("DB error: {e}"); std::process::exit(1); } }; let mut artists: Vec<(String, String, f64)> = vec![ (artist_mbid.to_string(), seed_name.to_string(), 1.0), ]; artists.extend(similar); let scored = playlist::score_tracks(conn, &artists, opts.popularity_bias); if opts.verbose { let mut sorted = scored.iter().collect::>(); sorted.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal)); for t in &sorted { eprintln!("{:.4}\t{:.4}\t{:.4}\t{}\t{}", t.score, t.similarity, t.popularity, t.artist, t.path); } } let candidates: Vec = scored .into_iter() .map(|t| Candidate { score: t.score, artist: t.artist, path: t.path, }) .collect(); let mut selected = playlist::generate_playlist(&candidates, opts.count, seed_name); if opts.random { selected.shuffle(&mut rand::rng()); } else if opts.shuffle { selected = playlist::interleave_artists(selected); } let tracks: Vec = selected.into_iter().map(|c| c.path).collect(); output_tracks(&tracks, opts, seed_name, conn); } fn output_tracks( tracks: &[String], opts: &BuildOptions, seed_name: &str, conn: &rusqlite::Connection, ) { if opts.mpd { let music_dir = env::var("MPD_MUSIC_DIR").unwrap_or_default(); if music_dir.is_empty() { eprintln!("Error: MPD_MUSIC_DIR not set"); std::process::exit(1); } let mut client = match mpd::MpdClient::connect() { Ok(c) => c, Err(e) => { eprintln!("MPD error: {e}"); std::process::exit(1); } }; client.queue_playlist(tracks, &music_dir); } else if opts.airsonic { let client = match airsonic::AirsonicClient::new() { Ok(c) => c, Err(e) => { eprintln!("Airsonic error: {e}"); std::process::exit(1); } }; if let Err(e) = client.create_playlist(seed_name, tracks, conn, opts.verbose) { eprintln!("Airsonic error: {e}"); std::process::exit(1); } } else { for track in tracks { println!("{track}"); } } }