drift/src/main.rs

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<String>,
    },
}

struct BuildOptions {
    verbose: bool,
    mpd: bool,
    airsonic: bool,
    shuffle: bool,
    random: bool,
    count: usize,
    popularity_bias: u8,
}

fn parse_positive_usize(s: &str) -> Result<usize, String> {
    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::<Vec<_>>();
        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<Candidate> = 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<String> = 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}");
        }
    }
}