selene_core/library/

extract.rs

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

use blake3::hash;
use lofty::{
    file::TaggedFileExt,
    picture::PictureType,
    tag::{ItemKey, TagItem, TagType},
};
use lunar_lib::database::{DatabaseEntry, writer::DatabaseWriter};

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

use crate::{
    database::{LibraryDb, Patchable, tx_extensions::CasTxExtensions},
    errors::ExtractError,
    library::{
        album::{Album, TrackReference, UNKNOWN_ALBUM},
        artist::{Artist, ArtistGroup, ArtistId},
        hash_source_files,
        metadata::{LoftyTagExtensions, extract_instrumental},
        track::{
            Track, TrackId, cover_art::CoverArt, lyric_data::LyricData, track_meta::TrackMeta,
        },
    },
    symphonia_helpers::{ContainerExtractResult, extract_from_file},
    utils::hash_file,
};

#[derive(Debug, Clone)]
pub struct ExtractResult {
    pub track: Track,
    pub album: Option<Album>,
    pub artists: Vec<Artist>,
}

/// Filters a slice of [`TrackId`]'s using the input
pub fn find_needs_extract(
    progress_renderer: Arc<dyn ProgressRenderer>,
) -> Result<Vec<PathBuf>, ExtractError> {
    let sources = hash_source_files(progress_renderer)?;
    let known_tracks = Track::db_get_all()?;
    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(
    files: &[PathBuf],
    progress_renderer: Arc<dyn ProgressRenderer>,
    dry: bool,
) -> Result<(), ExtractError> {
    let files: Vec<_> = files.iter().collect();

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

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

    let writer = DatabaseWriter::<LibraryDb>::spawn();

    files
        .par_iter()
        .try_for_each(|source| -> Result<(), ExtractError> {
            if writer.is_closed() {
                return Ok(());
            }

            let ExtractResult {
                track,
                album,
                artists,
            } = extract_metadata(source)?;

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

            if !dry {
                writer.transaction(move |cas_tx| {
                    cas_tx.tx_patch(track.clone())?;

                    if let Some(album) = &album {
                        cas_tx.tx_patch(album.clone())?;
                    }

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

                    Ok(())
                });
            }

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

            Ok(())
        })?;

    writer.finish()?;

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

/// Analyzes and extracts the input file metadata into an [`ExtractResult`]
pub fn extract_metadata(source_file: impl AsRef<Path>) -> Result<ExtractResult, 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();

        let has_album = has_album(
            title.as_deref(),
            album_title.as_deref(),
            &track_artists,
            &album_artists,
            track_num,
            track_total,
            disc_num,
            disc_total,
        );

        if has_album {
            let mut album = Album::new(
                album_title.unwrap_or(UNKNOWN_ALBUM.to_owned()),
                ArtistGroup::from_artists(&album_artists),
                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().patch(artist),
                    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: ArtistGroup::from_artists(&track_artists),
        date,
        genre,
        lyric_data: lyrics,
        other: other.collect(),
        title,
        art,
    };

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

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

    let mut all_artists: Vec<Artist> = all_artists.into_values().collect();

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

    let result = ExtractResult {
        track,
        album,
        artists: all_artists,
    };

    Ok(result)
}

fn has_album(
    track_title: Option<&str>,
    album_title: Option<&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 album_title.zip(track_title).is_none_or(|(a, b)| a != b) {
        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)
}