tag/src/tagger.rs

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> {
    // If the track already has an MBID, skip searching and go straight to lookup
    let (details, best_release) = if let Some(ref mbid) = track.musicbrainz_id {
        tracing::info!(id = track.id, mbid = %mbid, "looking up recording by MBID");

        if config.dry_run {
            tracing::info!(id = track.id, mbid = %mbid, "DRY RUN: would enrich from MBID");
            return Ok(true);
        }

        let details = provider.get_recording(mbid).await?;
        let best_release = details.releases.first().cloned();
        (details, best_release)
    } else {
        // No MBID — search by artist + title
        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");

        let candidates = provider.search_recording(&artist, &title).await?;

        if candidates.is_empty() {
            tracing::debug!(id = track.id, "no results from MusicBrainz");
            return Ok(false);
        }

        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);
        }

        let details = provider.get_recording(&best.recording.mbid).await?;
        let best_release = best.best_release;
        (details, best_release)
    };

    // Upsert primary 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 secondary collaborator artists so they exist as separate library entries
    for (name, mbid) in &details.secondary_artists {
        if let Err(e) = queries::artists::upsert(conn, name, Some(mbid)).await {
            tracing::warn!(artist = %name, error = %e, "failed to upsert secondary artist");
        }
    }

    // Upsert album from best release
    let (album_id, album_name) = if let Some(ref release) = 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_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
        && let Err(e) = file_tags::write_tags(
            &track.file_path,
            &details,
            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_needing_metadata(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)
}