Initial commit

src/tagger.rs

@@ -0,0 +1,227 @@
+use std::fmt;
+
+use sea_orm::{ActiveValue::Set, DatabaseConnection, NotSet};
+
+use shanty_db::entities::track;
+use shanty_db::queries;
+
+use crate::error::TagResult;
+use crate::file_tags;
+use crate::matcher::{self, ScoredMatch};
+use crate::provider::MetadataProvider;
+
+/// Configuration for a tagging operation.
+pub struct TagConfig {
+    /// If true, show what would change without writing to DB or files.
+    pub dry_run: bool,
+    /// If true, write updated tags back to the music files.
+    pub write_tags: bool,
+    /// Minimum match confidence (0.0 - 1.0).
+    pub confidence: f64,
+}
+
+/// Statistics from a completed tagging run.
+#[derive(Debug, Default, Clone)]
+pub struct TagStats {
+    pub tracks_processed: u64,
+    pub tracks_matched: u64,
+    pub tracks_updated: u64,
+    pub tracks_skipped: u64,
+    pub tracks_errored: u64,
+}
+
+impl fmt::Display for TagStats {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "processed: {}, matched: {}, updated: {}, skipped: {}, errors: {}",
+            self.tracks_processed,
+            self.tracks_matched,
+            self.tracks_updated,
+            self.tracks_skipped,
+            self.tracks_errored,
+        )
+    }
+}
+
+/// Tag a single track. Returns `Ok(true)` if matched and updated.
+pub async fn tag_track(
+    conn: &DatabaseConnection,
+    provider: &impl MetadataProvider,
+    track: &track::Model,
+    config: &TagConfig,
+) -> TagResult<bool> {
+    // Build search query
+    let (artist, title) = match matcher::build_query(track) {
+        Some(q) => q,
+        None => {
+            tracing::debug!(id = track.id, path = %track.file_path, "no query possible, skipping");
+            return Ok(false);
+        }
+    };
+
+    tracing::info!(
+        id = track.id,
+        artist = %artist,
+        title = %title,
+        "searching MusicBrainz"
+    );
+
+    // Search for recordings
+    let candidates = provider.search_recording(&artist, &title).await?;
+
+    if candidates.is_empty() {
+        tracing::debug!(id = track.id, "no results from MusicBrainz");
+        return Ok(false);
+    }
+
+    // Score and select best match
+    let best = match matcher::select_best_match(track, candidates, config.confidence) {
+        Some(m) => m,
+        None => {
+            tracing::debug!(
+                id = track.id,
+                "no match above confidence threshold {}",
+                config.confidence
+            );
+            return Ok(false);
+        }
+    };
+
+    log_match(track, &best);
+
+    if config.dry_run {
+        return Ok(true);
+    }
+
+    // Get full details for the best match
+    let details = provider.get_recording(&best.recording.mbid).await?;
+
+    // Upsert artist with MusicBrainz ID
+    let artist_id = match &details.artist_mbid {
+        Some(mbid) => {
+            Some(queries::artists::upsert(conn, &details.artist, Some(mbid)).await?.id)
+        }
+        None => {
+            Some(queries::artists::upsert(conn, &details.artist, None).await?.id)
+        }
+    };
+
+    // Upsert album from best release
+    let (album_id, album_name) = if let Some(ref release) = best.best_release {
+        let album = queries::albums::upsert(
+            conn,
+            &release.title,
+            &details.artist,
+            Some(&release.mbid),
+            artist_id,
+        )
+        .await?;
+        (Some(album.id), Some(release.title.clone()))
+    } else {
+        (None, None)
+    };
+
+    // Parse year from release date
+    let year = best
+        .best_release
+        .as_ref()
+        .and_then(|r| r.date.as_deref())
+        .and_then(|d| d.split('-').next())
+        .and_then(|y| y.parse::<i32>().ok());
+
+    let genre = details.genres.first().cloned();
+
+    // Update track metadata
+    let active = track::ActiveModel {
+        id: Set(track.id),
+        file_path: Set(track.file_path.clone()),
+        title: Set(Some(details.title.clone())),
+        artist: Set(Some(details.artist.clone())),
+        album: Set(album_name),
+        album_artist: Set(Some(details.artist.clone())),
+        musicbrainz_id: Set(Some(details.mbid.clone())),
+        artist_id: Set(artist_id),
+        album_id: Set(album_id),
+        year: Set(year),
+        genre: Set(genre.clone()),
+        // Preserve existing values for fields we don't update
+        track_number: NotSet,
+        disc_number: NotSet,
+        duration: NotSet,
+        codec: NotSet,
+        bitrate: NotSet,
+        file_size: NotSet,
+        fingerprint: NotSet,
+        file_mtime: NotSet,
+        added_at: NotSet,
+        updated_at: NotSet,
+    };
+    queries::tracks::update_metadata(conn, track.id, active).await?;
+
+    // Optionally write tags to file
+    if config.write_tags {
+        if let Err(e) = file_tags::write_tags(
+            &track.file_path,
+            &details,
+            best.best_release.as_ref(),
+            year,
+            genre.as_deref(),
+        ) {
+            tracing::warn!(id = track.id, path = %track.file_path, "failed to write file tags: {e}");
+        }
+    }
+
+    Ok(true)
+}
+
+fn log_match(track: &track::Model, best: &ScoredMatch) {
+    tracing::info!(
+        id = track.id,
+        confidence = format!("{:.2}", best.confidence),
+        matched_title = %best.recording.title,
+        matched_artist = %best.recording.artist,
+        release = best.best_release.as_ref().map(|r| r.title.as_str()).unwrap_or("(none)"),
+        "match found"
+    );
+}
+
+/// Run tagging on all untagged tracks or a specific track.
+pub async fn run_tagging(
+    conn: &DatabaseConnection,
+    provider: &impl MetadataProvider,
+    config: &TagConfig,
+    track_id: Option<i32>,
+) -> TagResult<TagStats> {
+    let tracks: Vec<track::Model> = if let Some(id) = track_id {
+        vec![queries::tracks::get_by_id(conn, id).await?]
+    } else {
+        queries::tracks::get_untagged(conn).await?
+    };
+
+    tracing::info!(count = tracks.len(), "tracks to process");
+    let mut stats = TagStats::default();
+
+    for track in &tracks {
+        stats.tracks_processed += 1;
+
+        match tag_track(conn, provider, track, config).await {
+            Ok(true) => {
+                stats.tracks_matched += 1;
+                if !config.dry_run {
+                    stats.tracks_updated += 1;
+                }
+            }
+            Ok(false) => {
+                stats.tracks_skipped += 1;
+            }
+            Err(e) => {
+                tracing::error!(id = track.id, path = %track.file_path, "tagging error: {e}");
+                stats.tracks_errored += 1;
+            }
+        }
+    }
+
+    tracing::info!(%stats, "tagging complete");
+    Ok(stats)
+}