SoulStream/src-tauri/src/loudness.rs

//! Volume normalization — the app-side orchestration over the player
//! crate's BS.1770 meter. Every main-player track is measured and played
//! with a gain that brings it to one target loudness (−14 LUFS, Spotify's
//! reference), so track-to-track volume stops being all over the place.
//!
//! Streaming-first: an unmeasured track is analyzed DURING its first play,
//! reading the same growing partial file the player streams from — analysis
//! runs at download speed, so the true integrated loudness usually lands
//! within seconds and the gain glides there (the player slew-limits it).
//! Measured tracks live in `loudness-cache.json`; their exact gain is
//! applied before the first sample on every later play.
//!
//! Toggleable via config `volume_normalization` (default on). Preview clips
//! are deliberately NOT normalized — Deezer masters them consistently.

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};

use serde::{Deserialize, Serialize};
use tauri::Manager;

use player::{analyze_loudness, AnalysisSource, StreamProgress};

use crate::playback::PlayerState;

/// Playback target (Spotify's default reference level).
pub const TARGET_LUFS: f64 = -14.0;
/// Gain clamp: quiet-track boosts are the risky direction (headroom), loud
/// tracks can always be cut.
const MIN_GAIN_DB: f64 = -15.0;
const MAX_GAIN_DB: f64 = 9.0;
/// Boosts additionally respect the track's own sample peak up to this
/// ceiling — normalization must never introduce clipping.
const PEAK_CEILING: f64 = 0.98;

const CACHE_FILE: &str = "loudness-cache.json";
/// Entries beyond this are dropped oldest-first at save time.
const CACHE_CAP: usize = 4000;

#[derive(Serialize, Deserialize, Clone, Copy)]
pub struct Measured {
    pub lufs: f64,
    pub peak: f32,
    /// File size when measured — a different size is a different file.
    pub size: u64,
    pub at_epoch_ms: i64,
}

/// Live view of the ACTIVE track's normalization, for the debug overlay.
#[derive(Serialize, Clone, Copy)]
#[serde(rename_all = "camelCase")]
pub struct CurrentTrack {
    pub lufs: Option<f64>,
    pub peak: f32,
    pub gain_db: f64,
    /// The gain came from the cache (exact, applied pre-sample).
    pub cached: bool,
    /// The analyzer is still measuring this play.
    pub analyzing: bool,
}

pub struct LoudnessState {
    cache: Mutex<HashMap<String, Measured>>,
    /// Bumped on every track start; orphans older analyzer threads.
    generation: Arc<AtomicU64>,
    enabled: AtomicBool,
    /// Debug-overlay mirror of the active track's measurement state.
    current: Mutex<Option<CurrentTrack>>,
}

/// (normalization enabled, target LUFS, active-track state) for diagnostics.
pub fn snapshot(app: &tauri::AppHandle) -> (bool, f64, Option<CurrentTrack>) {
    let state = app.state::<LoudnessState>();
    let current = *state.current.lock().expect("loudness current lock");
    (
        state.enabled.load(Ordering::Relaxed),
        TARGET_LUFS,
        current,
    )
}

fn cache_file(app: &tauri::AppHandle) -> Option<PathBuf> {
    app.path().app_data_dir().ok().map(|d| d.join(CACHE_FILE))
}

pub fn init(app: &tauri::AppHandle) {
    let cache: HashMap<String, Measured> = cache_file(app)
        .and_then(|p| std::fs::read_to_string(p).ok())
        .and_then(|s| serde_json::from_str(&s).ok())
        .unwrap_or_default();
    let enabled = crate::config::volume_normalization(&crate::config::load(app));
    app.manage(LoudnessState {
        cache: Mutex::new(cache),
        generation: Arc::new(AtomicU64::new(0)),
        enabled: AtomicBool::new(enabled),
        current: Mutex::new(None),
    });
}

fn save_cache(app: &tauri::AppHandle, state: &LoudnessState) {
    let Some(path) = cache_file(app) else { return };
    let mut entries: Vec<(String, Measured)> = state
        .cache
        .lock()
        .expect("loudness cache lock")
        .iter()
        .map(|(k, v)| (k.clone(), *v))
        .collect();
    if entries.len() > CACHE_CAP {
        entries.sort_by_key(|(_, m)| std::cmp::Reverse(m.at_epoch_ms));
        entries.truncate(CACHE_CAP);
    }
    let map: HashMap<&String, &Measured> = entries.iter().map(|(k, v)| (k, v)).collect();
    if let Ok(json) = serde_json::to_string(&map) {
        let _ = std::fs::write(path, json);
    }
}

/// gain = target − measured, clamped, and boosts capped so the track's own
/// peak never exceeds the ceiling.
fn gain_db(m: &Measured) -> f64 {
    let mut g = (TARGET_LUFS - m.lufs).clamp(MIN_GAIN_DB, MAX_GAIN_DB);
    if g > 0.0 && m.peak > 0.0 {
        let headroom = 20.0 * (PEAK_CEILING / m.peak as f64).log10();
        g = g.min(headroom.max(0.0));
    }
    g
}

/// Toggle from settings: applies to the CURRENT track immediately (gain to
/// unity when turning off; the cached gain back on when turning on is left
/// to the next track — the analyzer generation for the current one is gone).
pub fn set_enabled(app: &tauri::AppHandle, enabled: bool) {
    let state = app.state::<LoudnessState>();
    state.enabled.store(enabled, Ordering::Relaxed);
    if !enabled {
        state.generation.fetch_add(1, Ordering::Relaxed); // orphan analyzers
        *state.current.lock().expect("loudness current lock") = None;
        app.state::<PlayerState>().player.set_norm_gain_db(0.0, false);
    }
}

/// Pre-playback measurement window: the audible start is HELD until this
/// much audio is analyzed (at download/disk speed — normally well under a
/// second of wall time)…
const PRELIM_MIN_SECS: f64 = 10.0;
/// …or until this much wall time passes after the FIRST measurement lands
/// (slow peer): then we commit to what we have. Either way the gain is set
/// BEFORE the first audible sample and NEVER changes mid-track (user
/// decision 2026-07-23 — the earlier mid-play glide was audible and bad).
const PRELIM_WALL_BUDGET: std::time::Duration = std::time::Duration::from_millis(800);
/// A head estimate can under-read a track that opens quietly, so
/// preliminary BOOSTS are capped harder than cached ones; cuts are trusted
/// (a loud head is reliably loud).
const PRELIM_MAX_BOOST_DB: f64 = 3.0;

fn prelim_gain_db(m: &Measured) -> f64 {
    gain_db(m).min(PRELIM_MAX_BOOST_DB)
}

fn store_measurement(app: &tauri::AppHandle, key: &str, lufs: f64, peak: f32, size_hint: u64) {
    let final_size = std::fs::metadata(key).map(|m| m.len()).unwrap_or(size_hint);
    let measured = Measured {
        lufs,
        peak,
        size: final_size,
        at_epoch_ms: std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_millis() as i64)
            .unwrap_or(0),
    };
    let state = app.state::<LoudnessState>();
    state
        .cache
        .lock()
        .expect("loudness cache lock")
        .insert(key.to_owned(), measured);
    save_cache(app, &state);
}

/// The hook every main-player track start funnels through. `play` is the
/// DEFERRED start of playback: a cache hit or disabled normalization fires
/// it immediately; an unmeasured track holds it for the preliminary
/// measurement window so the gain is right from the very first sample.
/// The analyzer then keeps running silently to EOF — only to fill the
/// cache, never to touch the volume of the play in progress.
pub fn on_track_started(
    app: &tauri::AppHandle,
    cache_path: &Path,
    size: u64,
    stream: Option<(PathBuf, Arc<StreamProgress>)>,
    play: Box<dyn FnOnce() + Send>,
) {
    let state = app.state::<LoudnessState>();
    let generation = state.generation.fetch_add(1, Ordering::Relaxed) + 1;
    if !state.enabled.load(Ordering::Relaxed) {
        *state.current.lock().expect("loudness current lock") = None;
        play();
        return; // player resets to unity on Play
    }
    let key = cache_path.to_string_lossy().into_owned();

    // Known file → play at once with the exact gain (the set lands behind
    // Play on the ordered command channel, before any sink exists).
    let known_size = if size > 0 {
        size
    } else {
        std::fs::metadata(cache_path).map(|m| m.len()).unwrap_or(0)
    };
    if let Some(m) = state.cache.lock().expect("loudness cache lock").get(&key) {
        if m.size == known_size || known_size == 0 {
            let g = gain_db(m);
            play();
            app.state::<PlayerState>()
                .player
                .set_norm_gain_db(g as f32, true);
            *state.current.lock().expect("loudness current lock") = Some(CurrentTrack {
                lufs: Some(m.lufs),
                peak: m.peak,
                gain_db: g,
                cached: true,
                analyzing: false,
            });
            return;
        }
    }
    *state.current.lock().expect("loudness current lock") = Some(CurrentTrack {
        lufs: None,
        peak: 0.0,
        gain_db: 0.0,
        cached: false,
        analyzing: true,
    });

    // Unmeasured → the analyzer races ahead of playback (download/disk
    // speed) and RELEASES the held play once the preliminary window is met.
    let app = app.clone();
    let gen_ctr = state.generation.clone();
    std::thread::Builder::new()
        .name("loudness-analysis".into())
        .spawn(move || {
            let source = match stream {
                Some((partial, progress)) => AnalysisSource::Growing {
                    path: partial,
                    total: size,
                    progress,
                },
                None => AnalysisSource::Local(PathBuf::from(&key)),
            };
            let cancelled = || gen_ctr.load(Ordering::Relaxed) != generation;
            let mut play_slot = Some(play);
            let mut applied_db: Option<f64> = None;
            let mut first_measure: Option<std::time::Instant> = None;
            {
                let app = &app;
                let cancelled = &cancelled;
                let mut on_update = |lufs: f64, peak: f32, secs: f64| {
                    if cancelled() || play_slot.is_none() {
                        return; // started already: NO further volume touches
                    }
                    let first = *first_measure.get_or_insert_with(std::time::Instant::now);
                    if secs < PRELIM_MIN_SECS && first.elapsed() < PRELIM_WALL_BUDGET {
                        return; // keep listening a moment longer
                    }
                    let m = Measured {
                        lufs,
                        peak,
                        size: 0,
                        at_epoch_ms: 0,
                    };
                    let g = prelim_gain_db(&m);
                    (play_slot.take().expect("gate fires once"))();
                    app.state::<PlayerState>()
                        .player
                        .set_norm_gain_db(g as f32, true);
                    applied_db = Some(g);
                    *app.state::<LoudnessState>()
                        .current
                        .lock()
                        .expect("loudness current lock") = Some(CurrentTrack {
                        lufs: Some(lufs),
                        peak,
                        gain_db: g,
                        cached: false,
                        analyzing: true,
                    });
                };
                let done = analyze_loudness(source, cancelled, &mut on_update);

                // Gate never fired (clip shorter than a gating block, decode
                // failure, dead transfer): start at unity — unless a newer
                // track superseded this one, then never start it at all.
                if let Some(play) = play_slot.take() {
                    if !cancelled() {
                        play();
                    }
                }
                let Some((lufs, peak)) = done else { return };
                // The full measurement is valid for every future play of
                // this file, whatever happened to this one.
                store_measurement(app, &key, lufs, peak, size);
                if !cancelled() {
                    *app.state::<LoudnessState>()
                        .current
                        .lock()
                        .expect("loudness current lock") = Some(CurrentTrack {
                        lufs: Some(lufs),
                        peak,
                        gain_db: applied_db.unwrap_or(0.0),
                        cached: false,
                        analyzing: false,
                    });
                }
            }
        })
        .expect("spawn loudness analysis");
}

/// Gapless normalization: called when the NEXT track's preload is complete
/// (or was already cached on disk). Hands the boundary gain to the audio
/// thread, which applies it sample-accurately when the gapless switch
/// happens — a track change is the one place loudness may step. Guarded by
/// the loudness generation: a user-initiated play supersedes any pending
/// preload analysis (the player also clears its pending gain on Play/Stop).
pub fn on_preload_ready(app: &tauri::AppHandle, cache_path: &Path) {
    let state = app.state::<LoudnessState>();
    if !state.enabled.load(Ordering::Relaxed) {
        return;
    }
    let generation = state.generation.load(Ordering::Relaxed);
    let key = cache_path.to_string_lossy().into_owned();
    let known_size = std::fs::metadata(cache_path).map(|m| m.len()).unwrap_or(0);
    if let Some(m) = state.cache.lock().expect("loudness cache lock").get(&key) {
        if m.size == known_size || known_size == 0 {
            app.state::<PlayerState>()
                .player
                .set_next_norm_gain_db(gain_db(m) as f32);
            return;
        }
    }
    let app = app.clone();
    let gen_ctr = state.generation.clone();
    std::thread::Builder::new()
        .name("loudness-preload".into())
        .spawn(move || {
            let done = analyze_loudness(
                AnalysisSource::Local(PathBuf::from(&key)),
                &|| false,
                &mut |_, _, _| {},
            );
            let Some((lufs, peak)) = done else { return };
            store_measurement(&app, &key, lufs, peak, 0);
            // Only hand the boundary gain over if no new track started
            // since — a stale preload must not color a different queue.
            if gen_ctr.load(Ordering::Relaxed) == generation {
                let m = Measured {
                    lufs,
                    peak,
                    size: 0,
                    at_epoch_ms: 0,
                };
                app.state::<PlayerState>()
                    .player
                    .set_next_norm_gain_db(gain_db(&m) as f32);
            }
        })
        .expect("spawn preload loudness analysis");
}

#[cfg(test)]
mod tests {
    use super::*;

    fn m(lufs: f64, peak: f32) -> Measured {
        Measured {
            lufs,
            peak,
            size: 0,
            at_epoch_ms: 0,
        }
    }

    #[test]
    fn quiet_tracks_boost_toward_target() {
        assert!((gain_db(&m(-20.0, 0.3)) - 6.0).abs() < 1e-9);
    }

    #[test]
    fn loud_tracks_cut_toward_target() {
        assert!((gain_db(&m(-8.0, 1.0)) - (-6.0)).abs() < 1e-9);
    }

    #[test]
    fn boost_is_capped_by_peak_headroom() {
        // Target wants +9 (clamped), but a 0.9 peak only has ~0.73 dB of
        // headroom to the 0.98 ceiling.
        let g = gain_db(&m(-30.0, 0.9));
        assert!(g < 0.8 && g > 0.6, "got {g}");
    }

    #[test]
    fn cuts_ignore_peak() {
        assert!((gain_db(&m(-8.0, 0.2)) - (-6.0)).abs() < 1e-9);
    }

    #[test]
    fn extreme_measurements_stay_clamped() {
        assert_eq!(gain_db(&m(-60.0, 0.0001)), MAX_GAIN_DB);
        assert_eq!(gain_db(&m(20.0, 1.0)), MIN_GAIN_DB);
    }

    #[test]
    fn preliminary_boosts_are_capped_but_cuts_pass_through() {
        // A quiet head would ask +6 dB; the head estimate only gets +3.
        assert!((prelim_gain_db(&m(-20.0, 0.3)) - PRELIM_MAX_BOOST_DB).abs() < 1e-9);
        // Cuts are trusted at full strength — a loud head is reliably loud.
        assert!((prelim_gain_db(&m(-8.0, 1.0)) - (-6.0)).abs() < 1e-9);
    }
}