use std::{
collections::HashMap,
fs, io,
path::{Path, PathBuf},
sync::Arc,
};
use lunar_lib::{
database::{DatabaseEntry, DbIdExt, Entry, TransactionError, writer::DatabaseWriter},
id::Id,
iterator_ext::IteratorExtensions,
log::{debug, info, trace},
};
use rayon::iter::{IntoParallelRefIterator, ParallelBridge, ParallelIterator};
#[cfg(debug_assertions)]
use selene_core::database::validator::validate;
use selene_core::{
config::common_config,
database::{Library, Searchable},
library::{
album::Album,
artist::Artist,
extract_metadata, hash_source_files,
image_art::{ImageArt, Images},
track::{Track, lyric_data::LyricData},
},
};
use crate::{image_db, library_db};
pub fn extract() -> anyhow::Result<()> {
info!("Started extraction");
let can_multithread = common_config().main.multithreading;
let libary_db = library_db();
let image_db = image_db();
let sources_map = hash_source_files()?;
for entry in libary_db.iter_entries::<Track>() {
let mut entry = entry?;
let track_source = entry.container().path();
if track_source.try_exists()?
&& sources_map
.get(&entry.id())
.is_none_or(|t| t == track_source)
{
continue;
}
if let Some(source) = sources_map.get(&entry.id()) {
trace!("Relinked {}", entry.metadata.safe_title());
entry.__set_source_file(source.to_owned());
entry.db_upsert(libary_db)?;
} else {
entry.id().db_delete(libary_db)?;
}
}
#[cfg(debug_assertions)]
validate(libary_db)?;
let needs_extract = sources_map
.iter()
.filter_map(|(id, path)| match id.db_check(libary_db) {
Ok(false) => Some(Ok(path.to_owned())),
Ok(true) => None,
Err(e) => Some(Err(e)),
})
.try_to_vec()?;
if can_multithread {
let library_writer = DatabaseWriter::<Library>::spawn(libary_db.clone());
let image_writer = DatabaseWriter::<Images>::spawn(image_db.clone());
needs_extract
.par_iter()
.try_for_each(|source| -> anyhow::Result<()> {
if library_writer.is_closed() {
return Ok(());
}
let result = Arc::new(extract_metadata(source)?);
if library_writer.is_closed() {
return Ok(());
}
debug!("Extracted {}", result.track.metadata.safe_title());
let lib_result = result.clone();
library_writer
.transaction(move |cas_tx| (*lib_result).clone().tx_upsert_items(cas_tx));
image_writer.transaction(move |cas_tx| {
if let Some(image) = result.image()
&& !image.id().tx_check(cas_tx)?
{
Entry::from(image).tx_upsert(cas_tx)?;
}
Ok(())
});
Ok(())
})?;
library_writer.finish(true)?;
image_writer.finish(true)?;
} else {
needs_extract
.iter()
.try_for_each(|source| -> anyhow::Result<()> {
let result = extract_metadata(source)?;
debug!("Extracted {}", result.track.metadata.safe_title());
libary_db
.transaction(false, None, |cas_tx| {
result.clone().tx_upsert_items(cas_tx)?;
Ok(())
})
.map_err(TransactionError::from)?;
image_db
.transaction(false, None, |cas_tx| {
if let Some(image) = result.image()
&& !image.id().tx_check(cas_tx)?
{
Entry::from(image).tx_upsert(cas_tx)?;
}
Ok(())
})
.map_err(TransactionError::from)?;
Ok(())
})?;
}
#[cfg(debug_assertions)]
validate(libary_db)?;
for entry in libary_db.iter_entries::<Track>() {
let mut entry = entry?;
let mut lrc_file = entry.container().path().to_path_buf();
lrc_file.set_extension("slrc");
if !lrc_file.try_exists()? {
lrc_file.set_extension("lrc");
if !lrc_file.try_exists()? {
lrc_file.set_extension("txt");
if !lrc_file.try_exists()? {
continue;
}
}
}
let lyric_string = fs::read_to_string(lrc_file)?;
let Ok(lyric_data) = LyricData::infer_from_string(lyric_string) else {
continue;
};
trace!("Linked lyric data for {}", entry.metadata.safe_title());
entry.metadata.lyric_data = Some(lyric_data);
entry.db_upsert(libary_db)?;
}
let mut track_fs_groups: HashMap<PathBuf, Vec<_>> = HashMap::new();
for (track, mut parent) in sources_map {
let track = track
.db_get(libary_db)?
.expect("Extracted occured. All cached IDs should have database data");
assert!(parent.pop());
track_fs_groups.entry(parent).or_default().push(track);
}
let mut album_folder_map: HashMap<PathBuf, Id<Album>> = HashMap::new();
let mut value_counts: HashMap<Id<Album>, usize> = HashMap::new();
for (parent, group) in &track_fs_groups {
let Some(first_album) = group.iter().find_map(|track| track.metadata.album()) else {
continue;
};
if group
.iter()
.all(|track| track.metadata.album().is_none_or(|a| a == first_album))
{
album_folder_map.insert(parent.to_owned(), first_album);
*value_counts.entry(first_album).or_insert(0) += 1;
}
}
album_folder_map.retain(|_, album_id| value_counts[album_id] == 1);
for (directory, album_id) in album_folder_map {
let mut album = album_id
.db_get(libary_db)?
.expect("Album was verified to exist from previous tracks");
let mut group = track_fs_groups.remove(&directory).unwrap();
if album.art.is_none()
&& let Some(art) = art_from_dir(directory, group.len()).ok().flatten()
{
let id = art.id();
album.art = Some(id);
art.to_entry(id).db_upsert(image_db)?;
trace!("Linked album art for {}", album.name());
}
use selene_core::library::album::TrackReference;
for track in &mut group {
album.tracks.push(TrackReference {
id: track.id(),
track_num: None,
disc_num: None,
});
track.metadata.album = Some(album_id);
}
libary_db
.transaction(false, None, |cas_tx| {
for track in &group {
track.clone().tx_upsert(cas_tx)?;
}
album.clone().tx_upsert(cas_tx)?;
Ok(())
})
.map_err(TransactionError::from)?;
trace!("Linked album data for {}", album.name());
}
#[cfg(debug_assertions)]
validate(libary_db)?;
debug!("Building track search index");
Track::build_search_index(libary_db)?;
debug!("Building album search index");
Album::build_search_index(libary_db)?;
debug!("Building artist search index");
Artist::build_search_index(libary_db)?;
let accurate_true_peak = common_config().loudnorm.accurate_true_peak;
let track_iter = libary_db.iter_entries::<Track>();
if can_multithread {
let library_writer = DatabaseWriter::<Library>::spawn(libary_db.clone());
let max_threads = (rayon::current_num_threads() as f32 * 0.66).ceil() as usize;
let thread_pool = rayon::ThreadPoolBuilder::new()
.num_threads(max_threads)
.build()
.unwrap();
thread_pool.install(|| {
track_iter
.par_bridge()
.try_for_each(|track| -> anyhow::Result<()> {
if library_writer.is_closed() {
return Ok(());
}
let mut track = track?;
if track
.loudnorm_analysis()
.is_some_and(|la| la.accurate_true_peak() && accurate_true_peak)
{
return Ok(());
}
debug!("Extracting loudnorm for {}", track.metadata.safe_title());
track.loudnorm(accurate_true_peak)?;
if library_writer.is_closed() {
return Ok(());
}
library_writer.transaction(move |cas_tx| track.clone().tx_upsert(cas_tx));
Ok(())
})
})?;
} else {
for track in track_iter {
let mut track = track?;
if track
.loudnorm_analysis()
.is_some_and(|la| la.accurate_true_peak() && accurate_true_peak)
{
return Ok(());
}
debug!("Extracting loudnorm for {}", track.metadata.safe_title());
track.loudnorm(accurate_true_peak)?;
track.db_upsert(libary_db)?;
}
}
info!("Done extraction");
Ok(())
}
fn art_from_dir(dir: impl AsRef<Path>, expected_file_count: usize) -> io::Result<Option<ImageArt>> {
const COVER_LOOKUP_EXT: [&str; 3] = ["png", "jpeg", "jpg"];
let read_dir = fs::read_dir(dir.as_ref())?.flatten().to_vec();
if read_dir.len() != expected_file_count && read_dir.len() != expected_file_count + 1 {
return Ok(None);
}
let mut image = None;
for entry in read_dir {
if !entry.file_type().is_ok_and(|t| t.is_file()) {
continue;
}
let path = entry.path();
if path
.extension()
.and_then(std::ffi::OsStr::to_str)
.is_some_and(|ext| COVER_LOOKUP_EXT.contains(&ext))
&& let Ok(image_art) = ImageArt::from_image_file(path)
{
image = Some(image_art);
}
}
Ok(image)
}