selene-core 0.7.1

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

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

use crate::{
    config::common_config,
    database::{CasTxExtensions, LibraryDb, patch_vec},
    library::{
        album::{Album, TrackReference, UNKNOWN_ALBUM},
        artist::{Artist, ArtistId},
        hash_source_files,
        metadata::{LoftyTagExtensions, extract_instrumental},
        track::{
            Track, TrackId, cover_art::CoverArt, lyric_data::LyricData, track_meta::TrackMeta,
        },
    },
    media_container::ContainerError,
    symphonia_helpers::{ContainerExtractResult, extract_from_file},
    utils::hash_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,
}

#[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>,
    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(
    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...");

    #[inline]
    fn transaction_fn(
        cas_tx: &mut CompareAndSwapTransaction<LibraryDb>,
        track: &Track,
        album: Option<&Album>,
        artists: &[Artist],
    ) -> Result<(), TransactionError> {
        cas_tx.tx_upsert(track.clone())?;

        if let Some(album) = album {
            cas_tx.tx_merge(album.clone(), album.id())?;
        }

        for artist in artists {
            cas_tx.tx_merge(artist.clone(), artist.id())?;
        }

        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 ExtractResult {
                    track,
                    album,
                    artists,
                } = extract_metadata(source)?;

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

                if !dry {
                    writer.transaction(move |cas_tx| {
                        transaction_fn(cas_tx, &track, album.as_ref(), &artists)
                    });
                }

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

                Ok(())
            })?;
        writer.finish()?;
    } else {
        for source in files {
            let ExtractResult {
                track,
                album,
                artists,
            } = extract_metadata(source)?;

            if !dry {
                db_transaction(
                    |cas_tx| -> Result<(), CustomTransactionError<Infallible>> {
                        transaction_fn(cas_tx, &track, album.as_ref(), &artists)?;
                        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<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()),
                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) => {
                        patch_vec(&mut entry.get_mut().albums, 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(hash_file(source_file)?, 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: Vec<Artist> = all_artists.into_values().collect();

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