mtrack 0.12.0

// Copyright (C) 2026 Michael Wilson <mike@mdwn.dev>
//
// This program is free software: you can redistribute it and/or modify it under
// the terms of the GNU General Public License as published by the Free Software
// Foundation, version 3.
//
// This program is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
// FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License along with
// this program. If not, see <https://www.gnu.org/licenses/>.
//

//! Sample loading and caching for triggered samples.
//!
//! Samples are loaded entirely into memory at startup for zero-latency playback.

use std::collections::HashMap;
use std::error::Error;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;

use tracing::{debug, info, warn};

use crate::audio::sample_source::{create_sample_source_from_file, MemorySampleSource};
use crate::config::samples::SampleDefinition;

/// Default buffer size for reading samples (in samples, not bytes).
const DEFAULT_BUFFER_SIZE: usize = 4096;

/// A loaded sample that can be played back.
/// The sample data is stored in an Arc for efficient sharing between voices.
#[derive(Clone)]
pub struct LoadedSample {
    /// The sample data as f32 samples (interleaved if multi-channel).
    data: Arc<Vec<f32>>,
    /// Number of channels in the sample.
    channel_count: u16,
    /// Sample rate of the audio data.
    sample_rate: u32,
}

impl LoadedSample {
    /// Creates a new MemorySampleSource for playback with the given volume.
    pub fn create_source(&self, volume: f32) -> MemorySampleSource {
        MemorySampleSource::from_shared(
            self.data.clone(),
            self.channel_count,
            self.sample_rate,
            volume,
        )
    }

    /// Returns the number of channels.
    pub fn channel_count(&self) -> u16 {
        self.channel_count
    }

    /// Returns the memory size in bytes.
    pub fn memory_size(&self) -> usize {
        self.data.len() * std::mem::size_of::<f32>()
    }
}

/// Manages loading and caching of sample data.
pub struct SampleLoader {
    /// Cache of loaded samples by file path.
    cache: HashMap<PathBuf, LoadedSample>,
    /// Target sample rate for transcoding (matches audio output).
    target_sample_rate: u32,
}

impl SampleLoader {
    /// Creates a new sample loader.
    pub fn new(target_sample_rate: u32) -> Self {
        Self {
            cache: HashMap::new(),
            target_sample_rate,
        }
    }

    /// Loads a sample from a file into memory.
    /// Returns a cached version if already loaded.
    pub fn load(&mut self, path: &Path) -> Result<LoadedSample, Box<dyn Error>> {
        // Check cache first
        if let Some(sample) = self.cache.get(path) {
            debug!(path = ?path, "Using cached sample");
            return Ok(sample.clone());
        }

        info!(path = ?path, "Loading sample into memory");

        // Create a sample source from the file (include path in error)
        let mut source = create_sample_source_from_file(path, None, DEFAULT_BUFFER_SIZE).map_err(
            |e| -> Box<dyn std::error::Error> {
                format!("Failed to load sample {}: {}", path.display(), e).into()
            },
        )?;
        let source_sample_rate = source.sample_rate();
        let channel_count = source.channel_count();

        // Read all samples into memory
        let mut samples = Vec::new();
        while let Some(sample) = source.next_sample()? {
            samples.push(sample);
        }

        // Transcode if sample rate doesn't match
        let (final_samples, final_sample_rate) = if source_sample_rate != self.target_sample_rate {
            info!(
                source_rate = source_sample_rate,
                target_rate = self.target_sample_rate,
                "Transcoding sample"
            );
            let transcoded = self.transcode_samples(
                &samples,
                channel_count,
                source_sample_rate,
                self.target_sample_rate,
            )?;
            (transcoded, self.target_sample_rate)
        } else {
            (samples, source_sample_rate)
        };

        // Calculate duration
        let total_samples = final_samples.len();
        let samples_per_channel = total_samples as f64 / channel_count as f64;
        let duration_secs = samples_per_channel / final_sample_rate as f64;
        let duration = Duration::from_secs_f64(duration_secs);

        let loaded = LoadedSample {
            data: Arc::new(final_samples),
            channel_count,
            sample_rate: final_sample_rate,
        };

        info!(
            path = ?path,
            channels = channel_count,
            sample_rate = final_sample_rate,
            duration_ms = duration.as_millis(),
            memory_kb = loaded.memory_size() / 1024,
            "Sample loaded"
        );

        // Cache it
        self.cache.insert(path.to_path_buf(), loaded.clone());

        Ok(loaded)
    }

    /// Loads all samples referenced by a sample definition.
    /// Returns a map of file path to loaded sample.
    pub fn load_definition(
        &mut self,
        definition: &SampleDefinition,
        base_path: &Path,
    ) -> Result<HashMap<PathBuf, LoadedSample>, Box<dyn Error>> {
        let mut loaded = HashMap::new();

        for file in definition.all_files() {
            let full_path = if Path::new(file).is_absolute() {
                PathBuf::from(file)
            } else {
                base_path.join(file)
            };

            match self.load(&full_path) {
                Ok(sample) => {
                    loaded.insert(full_path, sample);
                }
                Err(e) => {
                    warn!(path = ?full_path, error = ?e, "Failed to load sample");
                    return Err(
                        format!("Failed to load sample {}: {}", full_path.display(), e).into(),
                    );
                }
            }
        }

        Ok(loaded)
    }

    /// Returns the total memory used by cached samples.
    pub fn total_memory_usage(&self) -> usize {
        self.cache.values().map(|s| s.memory_size()).sum()
    }

    /// Transcodes samples from one sample rate to another using linear interpolation.
    /// For higher quality, the existing Rubato transcoder could be used, but linear
    /// interpolation is simpler and often sufficient for drum hits and one-shots.
    fn transcode_samples(
        &self,
        samples: &[f32],
        channel_count: u16,
        source_rate: u32,
        target_rate: u32,
    ) -> Result<Vec<f32>, Box<dyn Error>> {
        let ratio = target_rate as f64 / source_rate as f64;
        let source_frames = samples.len() / channel_count as usize;
        let target_frames = (source_frames as f64 * ratio).ceil() as usize;
        let channels = channel_count as usize;

        let mut output = Vec::with_capacity(target_frames * channels);

        for target_frame in 0..target_frames {
            let source_pos = target_frame as f64 / ratio;
            let source_frame = source_pos.floor() as usize;
            let frac = source_pos.fract() as f32;

            for channel in 0..channels {
                let idx0 = source_frame * channels + channel;
                let idx1 = (source_frame + 1) * channels + channel;

                let s0 = samples.get(idx0).copied().unwrap_or(0.0);
                let s1 = samples.get(idx1).copied().unwrap_or(s0);

                // Linear interpolation
                let interpolated = s0 + (s1 - s0) * frac;
                output.push(interpolated);
            }
        }

        Ok(output)
    }
}

impl std::fmt::Debug for SampleLoader {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SampleLoader")
            .field("cached_samples", &self.cache.len())
            .field("target_sample_rate", &self.target_sample_rate)
            .field("total_memory_kb", &(self.total_memory_usage() / 1024))
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::samples::{VelocityConfig, VelocityLayer};

    #[test]
    fn test_transcode_samples() {
        let loader = SampleLoader::new(48000);

        // Simple mono sine wave at 44100Hz
        let source_rate = 44100;
        let target_rate = 48000;
        let source_samples: Vec<f32> = (0..4410)
            .map(|i| (2.0 * std::f32::consts::PI * 440.0 * i as f32 / source_rate as f32).sin())
            .collect();

        let result = loader
            .transcode_samples(&source_samples, 1, source_rate, target_rate)
            .unwrap();

        // Should have more samples at higher rate
        let expected_len = (4410.0_f64 * 48000.0 / 44100.0).ceil() as usize;
        assert_eq!(result.len(), expected_len);
    }

    #[test]
    fn test_transcode_stereo() {
        let loader = SampleLoader::new(48000);

        // Stereo: L=1.0, R=-1.0 alternating
        let source_samples = vec![1.0f32, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0];

        let result = loader
            .transcode_samples(&source_samples, 2, 44100, 48000)
            .unwrap();

        // Check that channels are preserved
        assert!(result.len() >= 8);
        // First frame should be close to original
        assert!((result[0] - 1.0).abs() < 0.1);
        assert!((result[1] - (-1.0)).abs() < 0.1);
    }

    #[test]
    fn test_transcode_same_rate() {
        let loader = SampleLoader::new(44100);
        let source = vec![0.5f32, -0.5, 0.25, -0.25];

        let result = loader.transcode_samples(&source, 1, 44100, 44100).unwrap();

        // Same rate: output length should match input length
        assert_eq!(result.len(), source.len());
        for (a, b) in result.iter().zip(source.iter()) {
            assert!((a - b).abs() < 1e-6);
        }
    }

    #[test]
    fn test_transcode_downsample() {
        let loader = SampleLoader::new(22050);
        let source: Vec<f32> = (0..4800).map(|i| (i as f32) / 4800.0).collect();

        let result = loader.transcode_samples(&source, 1, 48000, 22050).unwrap();

        // Downsampling should produce fewer samples
        assert!(result.len() < source.len());
        let expected_len = (4800.0_f64 * 22050.0 / 48000.0).ceil() as usize;
        assert_eq!(result.len(), expected_len);
    }

    #[test]
    fn test_transcode_empty_input() {
        let loader = SampleLoader::new(48000);
        let result = loader.transcode_samples(&[], 1, 44100, 48000).unwrap();
        assert!(result.is_empty());
    }

    #[test]
    fn test_transcode_single_frame() {
        let loader = SampleLoader::new(48000);
        let source = vec![0.75f32];

        let result = loader.transcode_samples(&source, 1, 44100, 48000).unwrap();

        // Should produce at least 1 sample
        assert!(!result.is_empty());
        // First sample should match source
        assert!((result[0] - 0.75).abs() < 0.01);
    }

    #[test]
    fn test_loaded_sample_memory_size() {
        let data = vec![1.0f32; 1000];
        let sample = LoadedSample {
            data: Arc::new(data),
            channel_count: 2,
            sample_rate: 44100,
        };

        assert_eq!(sample.memory_size(), 1000 * std::mem::size_of::<f32>());
        assert_eq!(sample.channel_count(), 2);
    }

    #[test]
    fn test_loaded_sample_create_source() {
        let data = vec![0.5f32; 100];
        let sample = LoadedSample {
            data: Arc::new(data),
            channel_count: 1,
            sample_rate: 44100,
        };

        // Should be able to create multiple sources from the same sample
        let _source1 = sample.create_source(1.0);
        let _source2 = sample.create_source(0.5);
    }

    #[test]
    fn test_sample_loader_new() {
        let loader = SampleLoader::new(48000);
        assert_eq!(loader.total_memory_usage(), 0);
    }

    #[test]
    fn test_sample_loader_debug() {
        let loader = SampleLoader::new(44100);
        let debug_str = format!("{:?}", loader);
        assert!(debug_str.contains("SampleLoader"));
        assert!(debug_str.contains("cached_samples"));
        assert!(debug_str.contains("target_sample_rate"));
    }

    #[test]
    fn test_load_wav_file() {
        let mut loader = SampleLoader::new(44100);
        let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("assets/1Channel44.1k.wav");

        let sample = loader.load(&path).unwrap();
        assert_eq!(sample.channel_count(), 1);
        assert!(sample.memory_size() > 0);
        assert!(loader.total_memory_usage() > 0);
    }

    #[test]
    fn test_load_caches_sample() {
        let mut loader = SampleLoader::new(44100);
        let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("assets/1Channel44.1k.wav");

        let sample1 = loader.load(&path).unwrap();
        let sample2 = loader.load(&path).unwrap();

        // Both should have the same data (Arc sharing)
        assert_eq!(sample1.memory_size(), sample2.memory_size());
        // Memory usage should not double
        assert_eq!(loader.total_memory_usage(), sample1.memory_size());
    }

    #[test]
    fn test_load_stereo_file() {
        let mut loader = SampleLoader::new(44100);
        let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("assets/2Channel44.1k.wav");

        let sample = loader.load(&path).unwrap();
        assert_eq!(sample.channel_count(), 2);
    }

    #[test]
    fn test_load_with_transcoding() {
        // Load a 22.05kHz file with a 44.1kHz target — should transcode
        let mut loader = SampleLoader::new(44100);
        let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("assets/1Channel22.05k.wav");

        let sample = loader.load(&path).unwrap();
        assert_eq!(sample.channel_count(), 1);
        assert!(sample.memory_size() > 0);
    }

    #[test]
    fn test_load_nonexistent_file() {
        let mut loader = SampleLoader::new(44100);
        let path = PathBuf::from("/nonexistent/file.wav");

        let result = loader.load(&path);
        assert!(result.is_err());
        assert!(result
            .err()
            .unwrap()
            .to_string()
            .contains("Failed to load sample"));
    }

    #[test]
    fn test_load_definition_single_file() {
        let mut loader = SampleLoader::new(44100);
        let base_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("assets");
        let definition = SampleDefinition::new(
            Some("1Channel44.1k.wav".to_string()),
            vec![1],
            VelocityConfig::ignore(None),
            crate::config::samples::ReleaseBehavior::PlayToCompletion,
            crate::config::samples::RetriggerBehavior::Cut,
            None,
            50,
        );

        let loaded = loader.load_definition(&definition, &base_path).unwrap();
        assert_eq!(loaded.len(), 1);
        assert!(loaded.contains_key(&base_path.join("1Channel44.1k.wav")));
    }

    #[test]
    fn test_load_definition_missing_file() {
        let mut loader = SampleLoader::new(44100);
        let base_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("assets");
        let definition = SampleDefinition::new(
            Some("nonexistent.wav".to_string()),
            vec![1],
            VelocityConfig::ignore(None),
            crate::config::samples::ReleaseBehavior::PlayToCompletion,
            crate::config::samples::RetriggerBehavior::Cut,
            None,
            50,
        );

        let result = loader.load_definition(&definition, &base_path);
        assert!(result.is_err());
    }

    #[test]
    fn test_load_definition_with_layers() {
        let mut loader = SampleLoader::new(44100);
        let base_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("assets");
        let layers = vec![
            VelocityLayer::new([1, 80], "1Channel44.1k.wav".to_string()),
            VelocityLayer::new([81, 127], "2Channel44.1k.wav".to_string()),
        ];
        let definition = SampleDefinition::new(
            None,
            vec![1],
            VelocityConfig::with_layers(layers, false),
            crate::config::samples::ReleaseBehavior::PlayToCompletion,
            crate::config::samples::RetriggerBehavior::Cut,
            None,
            50,
        );

        let loaded = loader.load_definition(&definition, &base_path).unwrap();
        assert_eq!(loaded.len(), 2);
    }

    #[test]
    fn test_transcode_preserves_stereo_channels() {
        let loader = SampleLoader::new(48000);

        // Create stereo signal: L channel positive ramp, R channel negative ramp
        let frames = 100;
        let mut source = Vec::with_capacity(frames * 2);
        for i in 0..frames {
            source.push(i as f32 / frames as f32); // L
            source.push(-(i as f32 / frames as f32)); // R
        }

        let result = loader.transcode_samples(&source, 2, 44100, 48000).unwrap();

        // Result should have even number of samples (stereo)
        assert_eq!(result.len() % 2, 0);

        // Check that L and R channels maintain their sign relationship
        for frame in result.chunks(2) {
            // L should be positive or zero, R should be negative or zero
            assert!(frame[0] >= -0.01, "L channel should be non-negative");
            assert!(frame[1] <= 0.01, "R channel should be non-positive");
        }
    }
}