prometeu-runtime/crates/host/prometeu-host-desktop-winit/src/audio.rs

use cpal::traits::{DeviceTrait, HostTrait, StreamTrait};
use prometeu_drivers::{AudioCommand, Channel, MAX_CHANNELS, OUTPUT_SAMPLE_RATE};
use prometeu_hal::LoopMode;
use ringbuf::traits::{Consumer, Producer, Split};
use ringbuf::HeapRb;
use std::error::Error;
use std::fmt;
use std::sync::Arc;
use std::time::Duration;

trait AudioStreamHandle {}

impl AudioStreamHandle for cpal::Stream {}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum HostAudioError {
    NoOutputDevice,
    BuildStream(String),
    PlayStream(String),
}

impl fmt::Display for HostAudioError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            HostAudioError::NoOutputDevice => write!(f, "no output device available"),
            HostAudioError::BuildStream(err) => write!(f, "failed to build audio stream: {}", err),
            HostAudioError::PlayStream(err) => write!(f, "failed to play audio stream: {}", err),
        }
    }
}

impl Error for HostAudioError {}

trait AudioRuntime {
    type Device;
    type Stream: AudioStreamHandle + 'static;

    fn default_output_device(&self) -> Option<Self::Device>;

    fn build_output_stream<D, E>(
        &self,
        device: &Self::Device,
        config: &cpal::StreamConfig,
        data_callback: D,
        error_callback: E,
    ) -> Result<Self::Stream, HostAudioError>
    where
        D: FnMut(&mut [f32]) + Send + 'static,
        E: FnMut(String) + Send + 'static;

    fn play_stream(&self, stream: &Self::Stream) -> Result<(), HostAudioError>;
}

struct CpalAudioRuntime;

impl AudioRuntime for CpalAudioRuntime {
    type Device = cpal::Device;
    type Stream = cpal::Stream;

    fn default_output_device(&self) -> Option<Self::Device> {
        cpal::default_host().default_output_device()
    }

    fn build_output_stream<D, E>(
        &self,
        device: &Self::Device,
        config: &cpal::StreamConfig,
        mut data_callback: D,
        mut error_callback: E,
    ) -> Result<Self::Stream, HostAudioError>
    where
        D: FnMut(&mut [f32]) + Send + 'static,
        E: FnMut(String) + Send + 'static,
    {
        device
            .build_output_stream(
                config,
                move |data: &mut [f32], _: &cpal::OutputCallbackInfo| data_callback(data),
                move |err| error_callback(err.to_string()),
                None,
            )
            .map_err(|err| HostAudioError::BuildStream(err.to_string()))
    }

    fn play_stream(&self, stream: &Self::Stream) -> Result<(), HostAudioError> {
        stream.play().map_err(|err| HostAudioError::PlayStream(err.to_string()))
    }
}

pub struct HostAudio {
    pub producer: Option<ringbuf::wrap::CachingProd<Arc<HeapRb<AudioCommand>>>>,
    pub perf_consumer: Option<ringbuf::wrap::CachingCons<Arc<HeapRb<u64>>>>,
    _stream: Option<Box<dyn AudioStreamHandle>>,
}

impl HostAudio {
    pub fn new() -> Self {
        Self { producer: None, perf_consumer: None, _stream: None }
    }

    pub fn init(&mut self) -> Result<(), HostAudioError> {
        self.init_with_runtime(&CpalAudioRuntime)
    }

    fn init_with_runtime<R: AudioRuntime>(&mut self, runtime: &R) -> Result<(), HostAudioError> {
        self.producer = None;
        self.perf_consumer = None;
        self._stream = None;

        let device = runtime.default_output_device().ok_or(HostAudioError::NoOutputDevice)?;

        let config = cpal::StreamConfig {
            channels: 2,
            sample_rate: cpal::SampleRate(OUTPUT_SAMPLE_RATE),
            buffer_size: cpal::BufferSize::Default,
        };

        let rb = HeapRb::<AudioCommand>::new(1024);
        let (prod, mut cons) = rb.split();

        let mut mixer = AudioMixer::new();

        // To pass performance data from the audio thread to the main thread
        let audio_perf_rb = HeapRb::<u64>::new(64);
        let (mut perf_prod, perf_cons) = audio_perf_rb.split();

        let stream = runtime.build_output_stream(
            &device,
            &config,
            move |data: &mut [f32]| {
                // Consumes commands from ringbuffer
                while let Some(cmd) = cons.try_pop() {
                    mixer.process_command(cmd);
                }
                // Mixes audio
                mixer.fill_buffer(data);
                // Sends processing time in microseconds
                let _ = perf_prod.try_push(mixer.last_processing_time.as_micros() as u64);
            },
            |err| eprintln!("audio stream error: {}", err),
        )?;

        runtime.play_stream(&stream)?;
        self.producer = Some(prod);
        self.perf_consumer = Some(perf_cons);
        self._stream = Some(Box::new(stream));
        Ok(())
    }

    pub fn send_commands(&mut self, commands: &mut Vec<AudioCommand>) {
        if let Some(producer) = &mut self.producer {
            for cmd in commands.drain(..) {
                if let Err(_) = producer.try_push(cmd) {
                    eprintln!("[HostAudio] Command ringbuffer full, dropping command.");
                }
            }
        } else {
            commands.clear();
        }
    }

    pub fn update_stats(&mut self, stats: &mut crate::stats::HostStats) {
        if let Some(cons) = &mut self.perf_consumer {
            while let Some(us) = cons.try_pop() {
                stats.record_audio_perf(us);
            }
        }
    }
}

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

    struct FakeStream;

    impl AudioStreamHandle for FakeStream {}

    #[derive(Clone, Copy)]
    struct FakeDevice;

    struct FakeRuntime {
        has_device: bool,
        build_result: Result<(), HostAudioError>,
        play_result: Result<(), HostAudioError>,
    }

    impl AudioRuntime for FakeRuntime {
        type Device = FakeDevice;
        type Stream = FakeStream;

        fn default_output_device(&self) -> Option<Self::Device> {
            self.has_device.then_some(FakeDevice)
        }

        fn build_output_stream<D, E>(
            &self,
            _device: &Self::Device,
            _config: &cpal::StreamConfig,
            _data_callback: D,
            _error_callback: E,
        ) -> Result<Self::Stream, HostAudioError>
        where
            D: FnMut(&mut [f32]) + Send + 'static,
            E: FnMut(String) + Send + 'static,
        {
            self.build_result.clone().map(|_| FakeStream)
        }

        fn play_stream(&self, _stream: &Self::Stream) -> Result<(), HostAudioError> {
            self.play_result.clone()
        }
    }

    #[test]
    fn test_send_commands_without_audio_drops_pending_commands() {
        let mut audio = HostAudio::new();
        let mut commands = vec![AudioCommand::MasterPause, AudioCommand::MasterResume];

        audio.send_commands(&mut commands);

        assert!(commands.is_empty());
        assert!(audio.producer.is_none());
    }

    #[test]
    fn test_update_stats_without_audio_is_safe() {
        let mut audio = HostAudio::new();
        let mut stats = HostStats::new();

        audio.update_stats(&mut stats);

        assert_eq!(stats.audio_load_accum_us, 0);
        assert_eq!(stats.audio_load_samples, 0);
    }

    #[test]
    fn test_init_returns_error_when_no_output_device_exists() {
        let mut audio = HostAudio::new();
        let runtime = FakeRuntime { has_device: false, build_result: Ok(()), play_result: Ok(()) };

        let result = audio.init_with_runtime(&runtime);

        assert_eq!(result, Err(HostAudioError::NoOutputDevice));
        assert!(audio.producer.is_none());
        assert!(audio.perf_consumer.is_none());
        assert!(audio._stream.is_none());
    }

    #[test]
    fn test_init_returns_error_when_stream_build_fails() {
        let mut audio = HostAudio::new();
        let runtime = FakeRuntime {
            has_device: true,
            build_result: Err(HostAudioError::BuildStream("simulated build failure".to_string())),
            play_result: Ok(()),
        };

        let result = audio.init_with_runtime(&runtime);

        assert_eq!(result, Err(HostAudioError::BuildStream("simulated build failure".to_string())));
        assert!(audio.producer.is_none());
        assert!(audio.perf_consumer.is_none());
        assert!(audio._stream.is_none());
    }

    #[test]
    fn test_init_returns_error_when_stream_play_fails() {
        let mut audio = HostAudio::new();
        let runtime = FakeRuntime {
            has_device: true,
            build_result: Ok(()),
            play_result: Err(HostAudioError::PlayStream("simulated play failure".to_string())),
        };

        let result = audio.init_with_runtime(&runtime);

        assert_eq!(result, Err(HostAudioError::PlayStream("simulated play failure".to_string())));
        assert!(audio.producer.is_none());
        assert!(audio.perf_consumer.is_none());
        assert!(audio._stream.is_none());
    }

    #[test]
    fn test_init_populates_audio_state_on_success() {
        let mut audio = HostAudio::new();
        let runtime = FakeRuntime { has_device: true, build_result: Ok(()), play_result: Ok(()) };

        audio.init_with_runtime(&runtime).unwrap();

        assert!(audio.producer.is_some());
        assert!(audio.perf_consumer.is_some());
        assert!(audio._stream.is_some());
    }
}

pub struct AudioMixer {
    voices: [Channel; MAX_CHANNELS],
    pub last_processing_time: Duration,
    paused: bool,
}

impl AudioMixer {
    pub fn new() -> Self {
        Self { voices: Default::default(), last_processing_time: Duration::ZERO, paused: false }
    }

    pub fn process_command(&mut self, cmd: AudioCommand) {
        match cmd {
            AudioCommand::Play { sample, voice_id, volume, pan, pitch, priority, loop_mode } => {
                if voice_id < MAX_CHANNELS {
                    println!(
                        "[AudioMixer] Playing voice {}: vol={}, pitch={}, loop={:?}",
                        voice_id, volume, pitch, loop_mode
                    );
                    self.voices[voice_id] = Channel {
                        sample: Some(sample),
                        active: true,
                        pos: 0.0,
                        pitch,
                        volume,
                        pan,
                        loop_mode,
                        priority,
                    };
                }
            }
            AudioCommand::Stop { voice_id } => {
                if voice_id < MAX_CHANNELS {
                    self.voices[voice_id].active = false;
                    self.voices[voice_id].sample = None;
                }
            }
            AudioCommand::SetVolume { voice_id, volume } => {
                if voice_id < MAX_CHANNELS {
                    self.voices[voice_id].volume = volume;
                }
            }
            AudioCommand::SetPan { voice_id, pan } => {
                if voice_id < MAX_CHANNELS {
                    self.voices[voice_id].pan = pan;
                }
            }
            AudioCommand::SetPitch { voice_id, pitch } => {
                if voice_id < MAX_CHANNELS {
                    self.voices[voice_id].pitch = pitch;
                }
            }
            AudioCommand::MasterPause => {
                println!("[AudioMixer] Master Pause");
                self.paused = true;
            }
            AudioCommand::MasterResume => {
                println!("[AudioMixer] Master Resume");
                self.paused = false;
            }
        }
    }

    pub fn fill_buffer(&mut self, buffer: &mut [f32]) {
        let start = std::time::Instant::now();
        // Zeroes the buffer (stereo)
        for sample in buffer.iter_mut() {
            *sample = 0.0;
        }

        if self.paused {
            self.last_processing_time = start.elapsed();
            return;
        }

        for voice in self.voices.iter_mut() {
            let sample_data = match &voice.sample {
                Some(s) => s,
                None => continue,
            };

            let pitch_ratio = sample_data.sample_rate as f64 / OUTPUT_SAMPLE_RATE as f64;
            let step = voice.pitch * pitch_ratio;

            let vol_f = voice.volume as f32 / 255.0;
            let pan_f = voice.pan as f32 / 255.0;
            let vol_l = vol_f * (1.0 - pan_f).sqrt();
            let vol_r = vol_f * pan_f.sqrt();

            for frame in buffer.chunks_exact_mut(2) {
                let pos_int = voice.pos as usize;
                let pos_fract = voice.pos - pos_int as f64;

                if pos_int >= sample_data.data.len() {
                    voice.active = false;
                    voice.sample = None;
                    break;
                }

                // Linear Interpolation
                let s1 = sample_data.data[pos_int] as f32 / 32768.0;
                let s2 = if pos_int + 1 < sample_data.data.len() {
                    sample_data.data[pos_int + 1] as f32 / 32768.0
                } else if voice.loop_mode == LoopMode::On {
                    let loop_start = sample_data.loop_start.unwrap_or(0) as usize;
                    sample_data.data[loop_start] as f32 / 32768.0
                } else {
                    0.0
                };

                let sample_val = s1 + (s2 - s1) * pos_fract as f32;

                frame[0] += sample_val * vol_l;
                frame[1] += sample_val * vol_r;

                voice.pos += step;

                let end_pos =
                    sample_data.loop_end.map(|e| e as f64).unwrap_or(sample_data.data.len() as f64);

                if voice.pos >= end_pos {
                    if voice.loop_mode == LoopMode::On {
                        let loop_start = sample_data.loop_start.unwrap_or(0) as f64;
                        voice.pos = loop_start + (voice.pos - end_pos);
                    } else {
                        voice.active = false;
                        voice.sample = None;
                        break;
                    }
                }
            }
        }

        // Final clamp to avoid clipping (optional if using f32, but good for fidelity)
        for sample in buffer.iter_mut() {
            *sample = sample.clamp(-1.0, 1.0);
        }
        self.last_processing_time = start.elapsed();
    }
}