selene_core/library/

extract.rs

use std::{
    collections::{HashMap, hash_map::Entry},
    convert::Infallible,
    path::{Path, PathBuf},
    sync::Arc,
};

use blake3::hash;
use lofty::{
    file::TaggedFileExt,
    picture::PictureType,
    tag::{ItemKey, TagItem, TagType},
};
use lunar_lib::{
    database::{
        CompareAndSwapTransaction, CustomTransactionError, DatabaseEntry, DbHandle,
        TransactionError, db_transaction, writer::DatabaseWriter,
    },
    vec_ext::VecExtensions,
};

use barber::{ProgressBar, ProgressRenderer};
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
use thiserror::Error;

use crate::{
    config::common_config,
    database::LibraryDb,
    library::{
        album::{Album, TrackReference},
        artist::{Artist, ArtistId},
        hash_source_files,
        metadata::{LoftyTagTakeAccessors, extract_instrumental},
        track::{
            ResolvedTrack, Track, TrackId, cover_art::CoverArt, lyric_data::LyricData,
            track_meta::TrackMeta,
        },
    },
    media_container::ContainerError,
    symphonia_helpers::{ContainerExtractResult, extract_from_file},
};

#[derive(Debug, Error)]
pub enum ExtractError {
    #[error("IoError: {0}")]
    Io(#[from] std::io::Error),

    #[error("LoftyError: {0}")]
    Lofty(#[from] lofty::error::LoftyError),

    #[error("DatabaseError: {0}")]
    Database(#[from] lunar_lib::database::DatabaseError),

    #[error("Transaction Error: {0}")]
    Transaction(#[from] TransactionError),

    #[error("Container Error: {0}")]
    Container(#[from] ContainerError),

    #[error("Template Error: {0}")]
    Template(#[from] lunar_lib::formatter::TemplateError),

    #[error("InvalidContainer Error")]
    InvalidContainer,
}

/// Filters a slice of [`TrackId`]'s using the input
pub fn find_needs_extract(
    progress_renderer: Arc<dyn ProgressRenderer>,
    db: &LibraryDb,
) -> Result<Vec<PathBuf>, ExtractError> {
    let sources = hash_source_files(progress_renderer)?;
    let known_tracks = Track::db_get_all(db)?;
    let known_track_ids: Vec<TrackId> = known_tracks.iter().map(Track::id).collect();

    Ok(sources
        .into_iter()
        .filter_map(|(id, path)| (!known_track_ids.contains(&id)).then_some(path))
        .collect())
}

pub fn extract(
    progress_renderer: Arc<dyn ProgressRenderer>,
    dry: bool,
) -> Result<(), ExtractError> {
    let db = DbHandle::<LibraryDb>::open()?;
    let files = find_needs_extract(progress_renderer.clone(), &db)?;

    if files.is_empty() {
        return Ok(());
    }

    let progress_bar = ProgressBar::new(0, files.len(), progress_renderer);
    progress_bar.set_label("Extracting metadata from files...");

    #[inline]
    fn transaction_fn(
        cas_tx: &mut CompareAndSwapTransaction<LibraryDb>,
        result: &ResolvedTrack,
    ) -> Result<(), TransactionError> {
        cas_tx.tx_upsert((*result.track).clone())?;

        if let Some((album, artists, ..)) = result.album_info() {
            cas_tx.tx_patch((**album).clone())?;
            for artist in artists {
                cas_tx.tx_patch((**artist).clone())?;
            }
        }

        for artist in result.artists() {
            cas_tx.tx_patch((**artist).clone())?;
        }

        Ok(())
    }

    if common_config().main.multithreading {
        let writer = DatabaseWriter::<LibraryDb>::spawn();
        files
            .par_iter()
            .try_for_each(|source| -> Result<(), ExtractError> {
                if writer.is_closed() {
                    return Ok(());
                }

                let result = extract_metadata(source)?;

                if writer.is_closed() {
                    return Ok(());
                }

                if !dry {
                    writer.transaction(move |cas_tx| transaction_fn(cas_tx, &result));
                }

                progress_bar.set_label(&format!(
                    "Extracted metadata from '{path}'",
                    path = source.display()
                ));
                progress_bar.increment();

                Ok(())
            })?;
        writer.finish()?;
    } else {
        for source in files {
            let result = extract_metadata(&source)?;

            if !dry {
                db_transaction(
                    |cas_tx| -> Result<(), CustomTransactionError<Infallible>> {
                        transaction_fn(cas_tx, &result)?;
                        Ok(())
                    },
                    DbHandle::<LibraryDb>::open().unwrap(),
                    false,
                )
                .map_err(TransactionError::from)?;
            }

            progress_bar.set_label(&format!(
                "Extracted metadata from '{path}'",
                path = source.display()
            ));
            progress_bar.increment();
        }
    }

    progress_bar.flush();
    Ok(())
}

/// Analyzes and extracts the input file metadata into an [`ExtractResult`]
pub fn extract_metadata(source_file: impl AsRef<Path>) -> Result<ResolvedTrack, ExtractError> {
    let source_file = source_file.as_ref();

    let ContainerExtractResult {
        container,
        // TODO: Replace lofty reading with symphonia
        metadata: _metadata,
    } = extract_from_file(source_file)?;

    let mut metadata = lofty::read_from_path(source_file)?;
    let tags = metadata.primary_tag_mut().unwrap();

    let (title, track_artists) = tags.track_title_and_artists();

    let (title, instrumental) = match title {
        Some(t) => {
            let (extracted, instrumental) = extract_instrumental(&t);
            (Some(extracted.into_owned()), Some(instrumental))
        }
        None => (None, None),
    };

    let mut all_artists: HashMap<ArtistId, Artist> =
        track_artists.iter().cloned().map(|a| (a.id(), a)).collect();

    let date = tags.date();

    let track_num = tags.track_num();
    let disc_num = tags.disc_num();

    let genre: Vec<String> = tags.take_strings(ItemKey::Genre).collect();

    let mut album = {
        let (album_title, album_artists) = tags.album_title_and_artists();
        let track_total = tags.track_total();
        let disc_total = tags.disc_total();

        if let Some(album_title) = album_title
            && has_album(
                title.as_deref(),
                &album_title,
                &track_artists,
                &album_artists,
                track_num,
                track_total,
                disc_num,
                disc_total,
            )
        {
            let mut album = Album::new(
                album_title,
                album_artists.iter().map(Artist::id).collect(),
                Vec::new(),
            );

            for mut artist in album_artists {
                artist.albums.push(album.id());

                match all_artists.entry(artist.id()) {
                    Entry::Occupied(mut entry) => {
                        entry.get_mut().albums.extend_unique(artist.albums);
                    }
                    Entry::Vacant(entry) => {
                        entry.insert(artist);
                    }
                }
            }

            album.date = date;
            album.track_total = track_total;
            album.disc_total = disc_total;
            album.genre = genre.clone();
            Some(album)
        } else {
            None
        }
    };

    let lyrics = tags.lyrics();
    let lyrics = if instrumental == Some(true) {
        Some(LyricData::Instrumental)
    } else {
        lyrics
    };

    let other = tags
        .items()
        .cloned()
        .map(TagItem::consume)
        .filter_map(|(k, v)| {
            let v = v.text()?.to_owned();
            let k = k.map_key(TagType::VorbisComments)?.to_owned();

            Some((k, v))
        });

    let art = tags
        .get_picture_type(PictureType::CoverFront)
        .or(tags.pictures().first())
        .map(|p| CoverArt::Embedded {
            hash: hash(p.data()),
            source: source_file.to_path_buf(),
        });

    let metadata = TrackMeta {
        album: album.as_ref().map(Album::id),
        artists: track_artists.iter().map(Artist::id).collect(),
        date,
        genre,
        lyric_data: lyrics,
        other: other.collect(),
        title,
        art,
    };

    let track = Track::new(container, metadata)?;

    if let Some(album) = &mut album {
        album.tracks.push(TrackReference {
            id: track.id(),
            track_num,
            disc_num,
        });
    }

    for artist in all_artists.values_mut() {
        if track_artists.contains(artist) {
            artist.tracks.push(track.id());
        }
    }

    let all_artists: HashMap<_, _> = all_artists
        .into_iter()
        .map(|(id, a)| (id, Arc::new(a)))
        .collect();

    let artists = track_artists
        .iter()
        .map(|a| all_artists[&a.id()].clone())
        .collect();

    let album_artists = album.as_ref().map(|album| {
        album
            .artists
            .iter()
            .map(|id| all_artists[id].clone())
            .collect()
    });

    Ok(ResolvedTrack {
        track: Arc::new(track),
        album: album.map(Arc::new),
        artists,
        album_artists,
    })
}

fn has_album(
    track_title: Option<&str>,
    album_title: &str,
    track_artists: &[Artist],
    album_artists: &[Artist],
    track_num: Option<u32>,
    track_total: Option<u32>,
    disc_num: Option<u32>,
    disc_total: Option<u32>,
) -> bool {
    if track_total == Some(1) && disc_total.is_none_or(|d| d <= 1) {
        return false;
    }

    // Returns true if the album name != track name
    if track_title.is_none_or(|t| t != album_title) {
        return true;
    }

    // Returns true if the album artists are not the same as the track artists
    if *album_artists != *track_artists {
        return true;
    }

    // Returns true if the track/disc value or total is greater is some and is greater than 1
    track_num.is_some_and(|v| v > 1)
        || track_total.is_some_and(|v| v > 1)
        || disc_num.is_some_and(|v| v > 1)
        || disc_total.is_some_and(|v| v > 1)
}