selene_core/library/

export.rs

use std::{
    collections::HashMap,
    fs,
    io::{self},
    path::{Path, PathBuf},
    sync::Arc,
};

use barber::{ProgressBar, ProgressRenderer};
use lofty::{config::WriteOptions, tag::TagExt};
use lunar_lib::{
    database::{DatabaseEntry, DbHandle, TransactionError},
    error,
    formatter::FormatTable,
};
use thiserror::Error;

use crate::{
    config::ExportConfig,
    database::{LibraryDb, Resolveable},
    library::{
        album::{Album, ResolvedAlbum},
        artist::add_from_artists,
        track::{ResolvedTrack, Track, lyric_data::LyricData},
    },
};

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

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

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

pub fn export_library(
    export_dir: impl AsRef<Path>,
    export_config: ExportConfig,
    progress_renderer: Arc<dyn ProgressRenderer>,
) -> Result<(), ExportError> {
    let db = DbHandle::<LibraryDb>::open().unwrap();
    let tracks = Track::db_get_all(&db)?;

    let mut conflicting_path_check: HashMap<PathBuf, usize> = HashMap::new();

    let mut track_export_targets = Vec::with_capacity(tracks.len());
    for track in tracks {
        let track = Track::resolve(Arc::new(track), &db)?;
        let export_path = {
            let mut format_table = FormatTable::new();
            format_table.extend_from_taggable(&track.metadata);

            add_from_artists(
                &mut format_table,
                track.artists().iter().map(|a| &**a),
                "track",
                &export_config.artist_separator,
                &export_config.alt_artist_separator,
            );

            if let Some((album, album_artists, track_num, disc_num)) = track.album_info() {
                format_table.extend_from_taggable(&**album);
                add_from_artists(
                    &mut format_table,
                    album_artists.iter().map(|a| &**a),
                    "album",
                    &export_config.artist_separator,
                    &export_config.alt_artist_separator,
                );
                if let Some(track_num) = track_num {
                    format_table.add_entry("track_num", track_num.to_string());
                }
                if let Some(disc_num) = disc_num {
                    format_table.add_entry("disc_num", disc_num.to_string());
                }
            }

            let path = format!(
                "{path}.{ext}",
                path = format_table.render(export_config.file_name_format.as_arguments()),
                ext = track.container().extension()
            );
            export_dir.as_ref().join(path)
        };

        if !export_config.overwrite && export_path.exists() {
            continue;
        }

        *conflicting_path_check
            .entry(export_path.clone())
            .or_default() += 1;
        track_export_targets.push((track, export_path));
    }

    let album_export_targets =
        if let Some(album_data_path_format) = export_config.album_data_path.as_ref() {
            let albums = Album::db_get_all(&db)?;
            let mut album_export_targets = Vec::with_capacity(albums.len());
            for album in albums {
                let album = Album::resolve(Arc::new(album), &db)?;
                let export_dir = {
                    let mut format_table = FormatTable::new();
                    format_table.extend_from_taggable(&*album);
                    add_from_artists(
                        &mut format_table,
                        album.artists.iter().map(|a| &**a),
                        "album",
                        &export_config.artist_separator,
                        &export_config.alt_artist_separator,
                    );

                    let path = format_table.render(album_data_path_format.as_arguments());
                    export_dir.as_ref().join(path)
                };

                *conflicting_path_check
                    .entry(export_dir.clone())
                    .or_default() += 1;
                album_export_targets.push((album, export_dir));
            }
            album_export_targets
        } else {
            Vec::new()
        };

    for (check, count) in conflicting_path_check {
        if count != 1 {
            error!("Conflicting path during export: {}", check.display());
            return Ok(());
        }
    }

    let progress_bar = ProgressBar::new(
        0,
        track_export_targets.len() + album_export_targets.len(),
        progress_renderer,
    );

    for (album, dir) in album_export_targets {
        export_album(&album, &export_config, dir)?;
        progress_bar.set_label(&format!("Exported data for album '{}'", album.name()));
        progress_bar.increment();
    }

    for (track, path) in track_export_targets {
        export_track(&track, &export_config, path)?;
        progress_bar.set_label(&format!(
            "Exported data for track '{}'",
            track.metadata().safe_title()
        ));
        progress_bar.increment();
    }

    Ok(())
}

pub fn export_album(
    album: &ResolvedAlbum,
    export_config: &ExportConfig,
    export_dir: impl AsRef<Path>,
) -> Result<(), ExportError> {
    let mut format_table = FormatTable::new();
    format_table.extend_from_taggable(&**album);
    add_from_artists(
        &mut format_table,
        album.artists.iter().map(|a| &**a),
        "album",
        &export_config.artist_separator,
        &export_config.alt_artist_separator,
    );

    // Cover art export
    if let Some(image_art) = album.art.as_ref()
        && export_config.album_data_path.is_some()
    {
        fs::create_dir_all(&export_dir)?;

        let ext = infer::get_from_path(image_art.source())?
            .map(|t| format!(".{}", t.extension()))
            .unwrap_or_default();

        let path = format!("cover{ext}");

        let cover_art_path = export_dir.as_ref().join(path);
        if export_config.overwrite || !cover_art_path.exists() {
            fs::copy(image_art.source(), cover_art_path)?;
        }
    }

    Ok(())
}

pub fn export_track(
    track: &ResolvedTrack,
    export_config: &ExportConfig,
    export_path: impl Into<PathBuf>,
) -> Result<(), ExportError> {
    let export_path = export_path.into();

    let buf = fs::read(track.container().path())?;
    let mut cursor = io::Cursor::new(buf);

    let mut tag = track.metadata_key_values(export_config);
    if let Some(cover_art) = track.metadata.art.as_ref() {
        tag.set_picture(0, cover_art.to_picture()?);
    }

    let options = WriteOptions::default().remove_others(true);
    tag.save_to(&mut cursor, options).unwrap();
    fs::create_dir_all(export_path.parent().unwrap())?;
    fs::write(&export_path, cursor.into_inner())?;

    if let Some(lyric_data) = track.metadata.lyric_data.as_ref() {
        match lyric_data {
            LyricData::Plain(plain_lyrics) if export_config.plain_lyrics_as_txt => {
                let mut lrc_path = export_path.clone();
                lrc_path.set_extension("txt");
                fs::write(lrc_path, plain_lyrics.as_str())?;
            }
            LyricData::Synced(synced_lyrics) => {
                if let Some((data, ext)) = export_config
                    .export_synced_lyrics_as
                    .map(|format| (synced_lyrics.to_lyrics(format), format.extension()))
                {
                    let mut lrc_path = export_path.clone();
                    lrc_path.set_extension(ext);

                    fs::write(lrc_path, data)?;
                }
            }
            _ => (),
        }
    }

    Ok(())
}