selene-core 0.8.2

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

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

use crate::{
    config::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(db: &LibraryDb) -> Result<Vec<Track>, DatabaseError> {
    let tracks = Track::db_get_all(db)?;
    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...");

    if common_config().main.multithreading {
        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.clone()));
                    }

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

                    Ok(())
                })
        })?;
        writer.finish()?;
    } else {
        for mut track in tracks {
            let analysis = LoudnormAnalysis::from_container(track.container())?;
            let title = track.metadata.safe_title().to_owned();

            track.loudnorm_analysis = Some(analysis);

            if !dry {
                db_transaction(
                    |cas_tx| cas_tx.tx_upsert(track.clone()).map_err(Into::into),
                    DbHandle::<LibraryDb>::open()?.clone(),
                    false,
                )
                .map_err(TransactionError::from)?;
            }

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

    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 {
    #[must_use]
    pub fn measured_i(&self) -> f64 {
        self.i
    }

    #[must_use]
    pub fn measured_tp(&self) -> f64 {
        self.tp
    }

    #[must_use]
    pub fn calculated_gain_db(&self) -> f64 {
        let config = common_config().loudnorm;
        (config.target_i - self.i - config.target_offset).min(config.target_tp - self.tp)
    }

    #[must_use]
    pub fn calculated_gain(&self) -> f64 {
        10f64.powf(self.calculated_gain_db() / 20.0)
    }

    #[must_use]
    pub fn calculated_replay_gain_peak(&self) -> f64 {
        10.0_f64.powf(self.tp / 20.0)
    }

    /// 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 {
                GenericAudioBufferRef::S16(buf) => {
                    let mut frames = vec![0; buf.samples_interleaved()];
                    buf.copy_to_slice_interleaved(&mut frames);
                    analyzer.add_frames_i16(&frames)?;
                }
                GenericAudioBufferRef::S32(buf) => {
                    let mut frames = vec![0; buf.samples_interleaved()];
                    buf.copy_to_slice_interleaved(&mut frames);
                    analyzer.add_frames_i32(&frames)?;
                }
                GenericAudioBufferRef::F64(buf) => {
                    let mut frames = vec![0.0; buf.samples_interleaved()];
                    buf.copy_to_slice_interleaved(&mut frames);
                    analyzer.add_frames_f64(&frames)?;
                }
                buf => {
                    let mut frames = vec![0.0; buf.samples_interleaved()];
                    buf.copy_to_slice_interleaved(&mut frames);
                    analyzer.add_frames_f32(&frames)?;
                }
            }
        }

        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 {
            accurate_true_peak,
            i,
            tp,
        })
    }
}