//! Master mix source: every layer of a song mixed frame-by-frame in a single
//! rodio source, with a soft-clip limiter on the sum, shared runtime volume
//! controls, sample-exact position and peak/RMS metering per layer and master.

use super::player::{SAMPLE_RATE, soft_clip};
use rodio::{ChannelCount, SampleRate, Source, source::SeekError};
use std::{
    num::NonZero,
    sync::{
        Arc,
        atomic::{AtomicU32, AtomicUsize, Ordering},
    },
    time::Duration,
};

const RATE: SampleRate = NonZero::new(SAMPLE_RATE).unwrap();
const STEREO: ChannelCount = NonZero::new(2).unwrap();
const METER_BLOCK_FRAMES: usize = 1024;
// Per-sample glide of applied layer gain toward its target, so volume changes
// (host `set_layer_volume`, parameter bindings) ramp instead of clicking.
// ponytail: fixed ≈5ms time constant at 44.1kHz; expose if a layer ever needs a
// custom fade rate.
const GAIN_GLIDE: f32 = 0.005;

/// Peak and RMS level of one mixer channel over the last metering block
/// (about 23 ms), post-fader. Values are linear amplitudes.
#[derive(Debug, Clone, Copy, Default)]
pub struct Levels {
    /// Highest absolute sample value in the block.
    pub peak: f32,
    /// Root-mean-square of the block.
    pub rms: f32,
}

struct Meter {
    peak: AtomicU32,
    rms: AtomicU32,
}

impl Meter {
    fn new() -> Self {
        Self {
            peak: AtomicU32::new(0),
            rms: AtomicU32::new(0),
        }
    }

    fn publish(&self, peak: f32, rms: f32) {
        self.peak.store(peak.to_bits(), Ordering::Relaxed);
        self.rms.store(rms.to_bits(), Ordering::Relaxed);
    }

    #[cfg(not(target_arch = "wasm32"))]
    fn levels(&self) -> Levels {
        Levels {
            peak: f32::from_bits(self.peak.load(Ordering::Relaxed)),
            rms: f32::from_bits(self.rms.load(Ordering::Relaxed)),
        }
    }
}

/// State shared between the control thread ([`super::SongPlayer`]) and the
/// [`MasterMix`] source on the audio thread.
pub(super) struct SharedMix {
    master_volume: AtomicU32,
    layer_volumes: Vec<AtomicU32>,
    layer_meters: Vec<Meter>,
    master_meter: Meter,
    position_frames: AtomicUsize,
}

impl SharedMix {
    pub fn new(layer_volumes: &[f32], master_volume: f32) -> Arc<Self> {
        Arc::new(Self {
            master_volume: AtomicU32::new(master_volume.to_bits()),
            layer_volumes: layer_volumes
                .iter()
                .map(|volume| AtomicU32::new(volume.to_bits()))
                .collect(),
            layer_meters: layer_volumes.iter().map(|_| Meter::new()).collect(),
            master_meter: Meter::new(),
            position_frames: AtomicUsize::new(0),
        })
    }

    #[cfg(not(target_arch = "wasm32"))]
    pub fn set_master_volume(&self, volume: f32) {
        self.master_volume
            .store(volume.to_bits(), Ordering::Relaxed);
    }

    fn master_volume(&self) -> f32 {
        f32::from_bits(self.master_volume.load(Ordering::Relaxed))
    }

    #[cfg(not(target_arch = "wasm32"))]
    pub fn set_layer_volume(&self, index: usize, volume: f32) {
        if let Some(slot) = self.layer_volumes.get(index) {
            slot.store(volume.to_bits(), Ordering::Relaxed);
        }
    }

    fn layer_volume(&self, index: usize) -> f32 {
        self.layer_volumes
            .get(index)
            .map_or(1.0, |slot| f32::from_bits(slot.load(Ordering::Relaxed)))
    }

    #[cfg(not(target_arch = "wasm32"))]
    pub fn layer_levels(&self, index: usize) -> Levels {
        self.layer_meters
            .get(index)
            .map_or_else(Levels::default, Meter::levels)
    }

    #[cfg(not(target_arch = "wasm32"))]
    pub fn master_levels(&self) -> Levels {
        self.master_meter.levels()
    }

    #[cfg(not(target_arch = "wasm32"))]
    pub fn position(&self) -> Duration {
        let frames = self.position_frames.load(Ordering::Relaxed);
        Duration::from_secs_f64(frames as f64 / f64::from(SAMPLE_RATE))
    }
}

/// One pre-rendered layer inside the mix: an interleaved stereo buffer with
/// an optional loop point.
pub(super) struct MixerLayer {
    samples: Vec<f32>,
    loop_start: Option<usize>,
    position: usize,
}

impl MixerLayer {
    /// Creates a layer starting `offset_frames` in (loop-aware: an offset past
    /// the end wraps into the loop region; one-shot layers just end).
    pub fn new(samples: Vec<f32>, loop_start_frame: Option<usize>, offset_frames: usize) -> Self {
        let loop_start = loop_start_frame.map(|frame| (frame * 2).min(samples.len()));
        let mut layer = Self {
            samples,
            loop_start,
            position: 0,
        };
        layer.seek_to_frame(offset_frames);
        layer
    }

    fn seek_to_frame(&mut self, frame: usize) {
        let mut sample = frame * 2;
        if let Some(loop_start) = self.loop_start
            && sample >= self.samples.len()
            && loop_start < self.samples.len()
        {
            let loop_length = self.samples.len() - loop_start;
            sample = loop_start + (sample - self.samples.len()) % loop_length;
        }
        self.position = sample.min(self.samples.len());
    }

    fn take_frame(&mut self) -> [f32; 2] {
        if self.position + 1 >= self.samples.len() {
            match self.loop_start {
                Some(loop_start) if loop_start + 1 < self.samples.len() => {
                    self.position = loop_start;
                }
                _ => return [0.0, 0.0],
            }
        }
        let frame = [self.samples[self.position], self.samples[self.position + 1]];
        self.position += 2;
        frame
    }

    fn finished(&self) -> bool {
        self.loop_start
            .is_none_or(|loop_start| loop_start + 1 >= self.samples.len())
            && self.position + 1 >= self.samples.len()
    }
}

struct Accumulator {
    peak: f32,
    sum_of_squares: f32,
}

impl Accumulator {
    fn new() -> Self {
        Self {
            peak: 0.0,
            sum_of_squares: 0.0,
        }
    }

    fn add(&mut self, frame: [f32; 2]) {
        for sample in frame {
            self.peak = self.peak.max(sample.abs());
            self.sum_of_squares += sample * sample;
        }
    }

    fn publish(&mut self, meter: &Meter, frames: usize) {
        let rms = (self.sum_of_squares / (frames as f32 * 2.0)).sqrt();
        meter.publish(self.peak, rms);
        self.peak = 0.0;
        self.sum_of_squares = 0.0;
    }
}

/// Wraps the whole mix back to `start_frame` once it reaches `end_frame`,
/// so a song that jumps to itself unconditionally loops gapless without a
/// re-render or sink swap.
pub(super) struct SongLoop {
    pub start_frame: usize,
    pub end_frame: usize,
}

/// The song's master source: mixes all layers, applies runtime layer and
/// master volumes, soft-clips the sum and feeds the shared meters.
pub(super) struct MasterMix {
    layers: Vec<MixerLayer>,
    shared: Arc<SharedMix>,
    current_gains: Vec<f32>,
    frame_count: usize,
    pending_right: Option<f32>,
    block_frames: usize,
    layer_accumulators: Vec<Accumulator>,
    master_accumulator: Accumulator,
    song_loop: Option<SongLoop>,
}

impl MasterMix {
    pub fn new(
        layers: Vec<MixerLayer>,
        shared: Arc<SharedMix>,
        offset_frames: usize,
        song_loop: Option<SongLoop>,
    ) -> Self {
        let layer_accumulators = layers.iter().map(|_| Accumulator::new()).collect();
        let current_gains = (0..layers.len())
            .map(|index| shared.layer_volume(index))
            .collect();
        shared
            .position_frames
            .store(offset_frames, Ordering::Relaxed);
        Self {
            layers,
            shared,
            current_gains,
            frame_count: offset_frames,
            pending_right: None,
            block_frames: 0,
            layer_accumulators,
            master_accumulator: Accumulator::new(),
            song_loop,
        }
    }
}

impl Iterator for MasterMix {
    type Item = f32;

    fn next(&mut self) -> Option<f32> {
        if let Some(right) = self.pending_right.take() {
            return Some(right);
        }
        if let Some(song_loop) = &self.song_loop {
            if self.frame_count >= song_loop.end_frame {
                for layer in &mut self.layers {
                    layer.seek_to_frame(song_loop.start_frame);
                }
                self.frame_count = song_loop.start_frame;
            }
        } else if self.layers.iter().all(MixerLayer::finished) {
            return None;
        }

        let mut mixed = [0.0f32; 2];
        for (index, layer) in self.layers.iter_mut().enumerate() {
            let target = self.shared.layer_volume(index);
            let gain = &mut self.current_gains[index];
            *gain += (target - *gain) * GAIN_GLIDE;
            let gain = *gain;
            let frame = layer.take_frame().map(|sample| sample * gain);
            self.layer_accumulators[index].add(frame);
            mixed[0] += frame[0];
            mixed[1] += frame[1];
        }
        let master = self.shared.master_volume();
        mixed = mixed.map(|sample| soft_clip(sample * master));
        self.master_accumulator.add(mixed);

        self.frame_count += 1;
        self.shared
            .position_frames
            .store(self.frame_count, Ordering::Relaxed);
        self.block_frames += 1;
        if self.block_frames >= METER_BLOCK_FRAMES {
            for (accumulator, meter) in self
                .layer_accumulators
                .iter_mut()
                .zip(&self.shared.layer_meters)
            {
                accumulator.publish(meter, self.block_frames);
            }
            self.master_accumulator
                .publish(&self.shared.master_meter, self.block_frames);
            self.block_frames = 0;
        }

        self.pending_right = Some(mixed[1]);
        Some(mixed[0])
    }
}

impl Source for MasterMix {
    fn current_span_len(&self) -> Option<usize> {
        None
    }

    fn channels(&self) -> ChannelCount {
        STEREO
    }

    fn sample_rate(&self) -> SampleRate {
        RATE
    }

    fn total_duration(&self) -> Option<Duration> {
        None
    }

    fn try_seek(&mut self, position: Duration) -> Result<(), SeekError> {
        let frame = (position.as_secs_f64() * f64::from(SAMPLE_RATE)) as usize;
        for layer in &mut self.layers {
            layer.seek_to_frame(frame);
        }
        self.frame_count = frame;
        self.shared.position_frames.store(frame, Ordering::Relaxed);
        self.pending_right = None;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    #![allow(clippy::unwrap_used)]
    use super::*;

    fn frames(mix: &mut MasterMix, count: usize) -> Vec<[f32; 2]> {
        (0..count)
            .map(|_| {
                let left = mix.next().unwrap_or(0.0);
                let right = mix.next().unwrap_or(0.0);
                [left, right]
            })
            .collect()
    }

    #[test]
    fn sums_layers_with_gains() {
        let shared = SharedMix::new(&[0.5, 1.0], 1.0);
        let layers = vec![
            MixerLayer::new(vec![0.4, 0.4, 0.4, 0.4], None, 0),
            MixerLayer::new(vec![0.2, 0.2], None, 0),
        ];
        let mut mix = MasterMix::new(layers, shared, 0, None);
        let output = frames(&mut mix, 2);
        assert!((output[0][0] - 0.4).abs() < 1e-6);
        assert!((output[1][0] - 0.2).abs() < 1e-6);
    }

    #[test]
    fn looping_layer_wraps_and_never_ends() {
        let shared = SharedMix::new(&[1.0], 1.0);
        let layers = vec![MixerLayer::new(vec![0.1, 0.1, 0.7, 0.7], Some(1), 0)];
        let mut mix = MasterMix::new(layers, shared, 0, None);
        let output = frames(&mut mix, 5);
        assert!((output[0][0] - 0.1).abs() < 1e-6);
        assert!((output[1][0] - 0.7).abs() < 1e-6);
        assert!((output[2][0] - 0.7).abs() < 1e-6);
        assert!((output[4][0] - 0.7).abs() < 1e-6);
    }

    #[test]
    fn one_shot_mix_ends() {
        let shared = SharedMix::new(&[1.0], 1.0);
        let layers = vec![MixerLayer::new(vec![0.5, 0.5], None, 0)];
        let mut mix = MasterMix::new(layers, shared, 0, None);
        assert!(mix.next().is_some());
        assert!(mix.next().is_some());
        assert!(mix.next().is_none());
    }

    #[test]
    fn master_sum_is_soft_clipped() {
        let shared = SharedMix::new(&[1.0, 1.0, 1.0], 1.0);
        let layers = (0..3)
            .map(|_| MixerLayer::new(vec![0.9, 0.9], None, 0))
            .collect();
        let mut mix = MasterMix::new(layers, shared, 0, None);
        let sample = mix.next().unwrap();
        assert!(sample <= 1.0);
        assert!(sample > 0.8);
    }

    #[test]
    fn position_and_meters_update() {
        let shared = SharedMix::new(&[1.0], 1.0);
        let samples = vec![0.5; METER_BLOCK_FRAMES * 2 + 4];
        let layers = vec![MixerLayer::new(samples, Some(0), 0)];
        let mut mix = MasterMix::new(layers, Arc::clone(&shared), 0, None);
        frames(&mut mix, METER_BLOCK_FRAMES + 1);
        assert_eq!(
            shared.position_frames.load(Ordering::Relaxed),
            METER_BLOCK_FRAMES + 1
        );
        let levels = shared.master_levels();
        assert!((levels.peak - 0.5).abs() < 1e-3);
        assert!((levels.rms - 0.5).abs() < 1e-3);
        assert!((shared.layer_levels(0).peak - 0.5).abs() < 1e-3);
    }

    #[test]
    fn seek_wraps_looping_layers() {
        let shared = SharedMix::new(&[1.0], 1.0);
        let layers = vec![MixerLayer::new(vec![0.1, 0.1, 0.7, 0.7], Some(1), 0)];
        let mut mix = MasterMix::new(layers, shared, 0, None);
        mix.try_seek(Duration::from_secs_f64(5.0 / f64::from(SAMPLE_RATE)))
            .unwrap();
        let output = frames(&mut mix, 1);
        assert!((output[0][0] - 0.7).abs() < 1e-6);
    }

    #[test]
    fn song_loop_wraps_the_whole_mix() {
        let shared = SharedMix::new(&[1.0], 1.0);
        let layers = vec![MixerLayer::new(vec![0.1, 0.1, 0.2, 0.2, 0.3, 0.3], None, 0)];
        let song_loop = SongLoop {
            start_frame: 1,
            end_frame: 3,
        };
        let mut mix = MasterMix::new(layers, shared, 0, Some(song_loop));
        let output = frames(&mut mix, 5);
        assert!((output[0][0] - 0.1).abs() < 1e-6);
        assert!((output[2][0] - 0.3).abs() < 1e-6);
        assert!((output[3][0] - 0.2).abs() < 1e-6);
        assert!((output[4][0] - 0.3).abs() < 1e-6);
    }
}
