selene-core 0.5.0

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 symphonia::core::audio::{Audio, AudioBuffer, GenericAudioBuffer};
use thiserror::Error;

use crate::{
    config::common::common_config,
    database::LibraryDb,
    library::track::Track,
    media_container::MediaContainer,
    symphonia_helpers::raw_decoder::{DecodingError, RawDecoder},
};

/// 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()?;
    let accurate_true_peak = common_config().loudnorm.accurate_true_peak;

    Ok(tracks
        .into_iter()
        .filter(|t| {
            t.loudnorm_analysis
                .is_none_or(|a| !a.accurate_true_peak && accurate_true_peak)
        })
        .collect())
}

pub fn analyze_loudorm(
    tracks: Vec<Track>,
    progress_renderer: Arc<dyn ProgressRenderer>,
    dry: bool,
) -> 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.container())?;

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

                track.loudnorm_analysis = Some(analysis);

                let title = track.metadata.safe_title().to_owned();

                if !dry {
                    writer.transaction(move |cas_tx| {
                        cas_tx.tx_upsert(track.id(), Some(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) accurate_true_peak: bool,

    pub(crate) i: f64,
    pub(crate) tp: f64,
}

impl LoudnormAnalysis {
    pub fn measured_i(&self) -> f64 {
        self.i
    }

    pub fn measured_tp(&self) -> f64 {
        self.tp
    }

    /// Uses ffmpeg to measure the input file and return a [`LoudnormAnalysis`]
    pub fn from_container(container: &MediaContainer) -> Result<Self, LoudnormAnalysisError> {
        let accurate_true_peak = common_config().loudnorm.accurate_true_peak;

        let peak_mode = if accurate_true_peak {
            Mode::TRUE_PEAK
        } else {
            Mode::SAMPLE_PEAK
        };

        let mut analyzer = EbuR128::new(
            container.stream().codec_params.channels as u32,
            container.stream().codec_params.sample_rate,
            Mode::I | peak_mode,
        )?;

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

        while let Some(packet) = raw_decoder.decode_next_packet()? {
            match packet {
                GenericAudioBuffer::S16(buf) => {
                    analyzer.add_frames_i16(&buf.iter_interleaved().collect::<Vec<i16>>())?
                }
                GenericAudioBuffer::S32(buf) => {
                    analyzer.add_frames_i32(&buf.iter_interleaved().collect::<Vec<i32>>())?
                }
                GenericAudioBuffer::F32(buf) => {
                    analyzer.add_frames_f32(&buf.iter_interleaved().collect::<Vec<f32>>())?
                }
                GenericAudioBuffer::F64(buf) => {
                    analyzer.add_frames_f64(&buf.iter_interleaved().collect::<Vec<f64>>())?
                }
                other => {
                    let mut buf = AudioBuffer::<f32>::new(other.spec().clone(), other.frames());
                    buf.render_silence(Some(other.frames()));
                    other.copy_to(&mut buf);
                    analyzer.add_frames_f32(&buf.iter_interleaved().collect::<Vec<f32>>())?
                }
            }
        }

        let i = analyzer.loudness_global()?;
        let tp = (0..container.stream().codec_params.channels as u32).try_fold(
            f64::NEG_INFINITY,
            |max, c| {
                let peak = if accurate_true_peak {
                    analyzer.true_peak(c)
                } else {
                    analyzer.sample_peak(c)
                };
                peak.map(|tp| f64::max(max, tp))
            },
        )?;

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