selene-core 0.4.2

use std::{f64, sync::Arc};

use barber::{ProgressBar, ProgressRenderer};
use ebur128::{EbuR128, Mode};
use lunar_lib::database::{DatabaseEntry, DatabaseError, TransactionError, writer::DatabaseWriter};
use rayon::{
    ThreadPoolBuilder,
    iter::{IntoParallelIterator, ParallelIterator},
};
use serde::{Deserialize, Serialize};
use thiserror::Error;

use crate::{
    database::{LibraryDb, tx_extensions::CasTxExtensions},
    decoding::{DecodingError, RawDecoder},
    library::track::Track,
    media_container::MediaContainer,
};

/// Filters a slice of [`TrackId`]'s using the input
pub fn find_needs_loudnorm_analysis() -> Result<Vec<Track>, DatabaseError> {
    let tracks = Track::db_get_all()?;

    Ok(tracks
        .into_iter()
        .filter(|t| t.loudnorm_analysis.is_none())
        .collect())
}

pub fn analyze_loudorm(
    tracks: Vec<Track>,
    progress_renderer: Arc<dyn ProgressRenderer>,
) -> Result<(), LoudnormAnalysisError> {
    if tracks.is_empty() {
        return Ok(());
    }

    let progress_bar = ProgressBar::new(0, tracks.len(), progress_renderer);
    progress_bar.set_label("Extracting loudnorm analysis from source files...");

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

    let max_threads = (rayon::current_num_threads() as f32 * 0.66).ceil() as usize;
    let thread_pool = ThreadPoolBuilder::new()
        .num_threads(max_threads)
        .build()
        .unwrap();

    thread_pool.install(|| {
        tracks
            .into_par_iter()
            .try_for_each(|mut track| -> Result<(), LoudnormAnalysisError> {
                if writer.is_closed() {
                    return Ok(());
                }

                let analysis = LoudnormAnalysis::from_container(track.src_container())?;

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

                track.loudnorm_analysis = Some(analysis);

                let title = track.metadata.safe_title().to_owned();
                writer.transaction(move |cas_tx| cas_tx.tx_patch(track.clone()));

                progress_bar.set_label(&format!("Extracted loudnorm for '{title}'",));
                progress_bar.increment();

                Ok(())
            })
    })?;

    writer.finish()?;

    progress_bar.flush();

    Ok(())
}

#[derive(Debug, Error)]
pub enum LoudnormAnalysisError {
    #[error("{0}")]
    Ebur128(#[from] ebur128::Error),

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

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

    #[error("Decoding Error: {0}")]
    Decoding(#[from] DecodingError),
}

/// [`LoudnormAnalysis`] represents the measured values from the first-pass of EBU R 128 loudnorm
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, PartialOrd, Copy)]
pub struct LoudnormAnalysis {
    pub(crate) i: f64,
    pub(crate) tp: f64,
    pub(crate) lra: f64,
}

impl LoudnormAnalysis {
    /// Uses ffmpeg to measure the input file and return a [`LoudnormAnalysis`]
    pub fn from_container(container: &MediaContainer) -> Result<Self, LoudnormAnalysisError> {
        let mut analyzer = EbuR128::new(
            container.stream().channels as u32,
            container.stream().sample_rate,
            Mode::I | Mode::LRA | Mode::TRUE_PEAK,
        )?;

        let mut raw_decoder = RawDecoder::from_container(container, 512 * 1024)?;

        while let Some(packet) = raw_decoder.decode_next_packet()? {
            let samples = packet.sample_buf.samples();
            analyzer.add_frames_f32(samples)?;
        }

        let i = analyzer.loudness_global()?;
        let lra = analyzer.loudness_range()?;
        let tp = (0..container.stream().channels)
            .map(|c| analyzer.true_peak(c as u32))
            .collect::<Result<Vec<_>, _>>()?
            .into_iter()
            .fold(f64::NEG_INFINITY, f64::max);

        Ok(LoudnormAnalysis { i, tp, lra })
    }
}