eva-pwm-cohost/src/sfx.rs

use std::path::Path;

use rand::seq::IteratorRandom;
use resampler::ResamplerFir;
use symphonia::core::{codecs::audio::AudioDecoder, formats::{FormatReader, TrackType, probe::Hint}, io::MediaSourceStream};

use crate::audio::AudioOutStream;

#[derive(Debug)]
pub enum SfxRequest {
    RandomAmbient
}

#[derive(Debug, Clone)]
pub struct SfxControl {
    sink: tokio::sync::mpsc::Sender<SfxRequest>
}

impl SfxControl {
    pub async fn play_ambient(&mut self) -> Result<(), tokio::sync::mpsc::error::SendError<SfxRequest>> {
        self.sink.send(SfxRequest::RandomAmbient).await
    }
}

struct Sample {
    format: Box<dyn FormatReader>,
    decoder: Box<dyn AudioDecoder>,
    channel_bufs: Vec<f32>,
    sample_rate: u32,
    channel_num: usize,
    bitrate_resample: ResamplerFir
}

impl Sample {
    fn new(format: Box<dyn FormatReader>, decoder: Box<dyn AudioDecoder>, output_sample_rate: u32) -> Self {
        let sample_rate = decoder.codec_params().sample_rate.unwrap();
        let channel_num = decoder.codec_params().channels.as_ref().unwrap().count();
        let bitrate_resample = resampler::ResamplerFir::new_from_hz(channel_num, sample_rate, output_sample_rate, Default::default(), Default::default());
        Self {
            format,
            decoder,
            channel_bufs: vec![],
            sample_rate,
            channel_num,
            bitrate_resample
        }
    }

    fn next(&mut self) -> Result<Option<Vec<f32>>, symphonia::core::errors::Error> {
        let packet = match self.format.next_packet()? {
            Some(packet) => packet,
            None => return Ok(None)
        };

        match self.decoder.decode_ref(&packet.as_packet_ref()) {
            Ok(samples) => {
                self.channel_bufs.resize(samples.samples_interleaved(), 0.);
                samples.copy_to_slice_interleaved(&mut self.channel_bufs);

                let mut resampled = [0.; 4096];
                let (read_count, write_count) = self.bitrate_resample.resample(&self.channel_bufs, &mut resampled).unwrap();
                if read_count < self.channel_bufs.len() {
                    log::error!("Resampling buffer is too small for a {}Hz file! We need an additional {}", self.sample_rate, read_count);
                }

                // First we convert the audio feed from stereo down to mono by simple average
                // TODO: This should be something smarter, like a saturating add..?
                let mono_stream = resampled[..write_count].chunks(self.channel_num).map(|channels| {
                    let total_volume = channels.iter().cloned().reduce(|a, b| a + b).unwrap_or_default();
                    total_volume / (self.channel_num as f32)
                });

                Ok(Some(mono_stream.collect()))
            },
            Err(err) => {
                // Dump the audio buffer on failure
                Err(err)
            }
        }
    }
}

struct Player {
    audio_sink: AudioOutStream,
    audio_out_buf: Vec<f32>,
    playing_samples: Vec<Sample>
}

impl Player {
    async fn process(&mut self) {
        if self.playing_samples.is_empty() {
            return;
        }

        'out: loop {
            let mut next_batch = vec![];

            self.playing_samples.retain_mut(|sample| {
                match sample.next() {
                    Ok(Some(buf)) => {
                        next_batch.push(buf);
                        true
                    },
                    Ok(None) => {
                        false
                    },
                    Err(err) => {
                        log::error!("Audio error: {:?}", err);
                        false
                    }
                }
            });

            loop {
                let mut this_sample = 0.;
                let mut any_valid = false;
                for stream in &mut next_batch {
                    if let Some(next_sample) = stream.pop() {
                        this_sample += next_sample;
                        any_valid = true;
                    }
                }

                if !any_valid {
                    if self.audio_out_buf.is_empty() {
                        break 'out;
                    }
                    break;
                }
                let mixed_sample = this_sample / next_batch.len() as f32;
                self.audio_out_buf.push(mixed_sample);
                if self.audio_out_buf.len() >= 1024 {
                    self.flush().await;
                }
            }

            if self.playing_samples.is_empty() {
                break;
            }
        }
    }

    async fn submit_buffer(&mut self) {
        self.audio_sink.sink.send(std::mem::take(&mut self.audio_out_buf)).await.unwrap();
    }

    async fn flush(&mut self) {
        if !self.audio_out_buf.is_empty() {
            self.submit_buffer().await;
        }
    }

    async fn play_sample(&mut self, path: &Path) {
        log::debug!("Queuing sound playback for {:?}", path);
        let sfx_fd = std::fs::File::open(path).unwrap();
        let mss = MediaSourceStream::new(Box::new(sfx_fd), Default::default());
        let meta_opts = Default::default();
        let fmt_opts = Default::default();
        let mut hint = Hint::new();
        // FIXME: use actual file extension
        hint.with_extension(".mp3");
        let format = symphonia::default::get_probe()
            .probe(&hint, mss, fmt_opts, meta_opts)
            .expect("Unsupported audio format");
        let track = format.default_track(TrackType::Audio).expect("No audio track");
        let dec_opts = Default::default();
        let decoder = symphonia::default::get_codecs()
            .make_audio_decoder(
                track.codec_params.as_ref().expect("codec params missing").audio().unwrap(),
                &dec_opts
            ).expect("Unsupported audio codec");

        self.playing_samples.push(Sample::new(format, decoder, self.audio_sink.sample_rate));
    }
}

pub async fn start_sfx(audio_sink: AudioOutStream) -> SfxControl {
    let (event_sink, mut event_src) = tokio::sync::mpsc::channel(32);
    tokio::spawn(async move {
        let mut player = Player { audio_sink, audio_out_buf: vec![] , playing_samples: vec![]};
        let sfx_dir = std::path::Path::new("./sfx");

        loop {
            tokio::select! {
                _ = player.process(), if !player.playing_samples.is_empty() => {},
                Some(event) = event_src.recv() => {
                    match event {
                        SfxRequest::RandomAmbient => {
                            log::debug!("Playing random audio sample");
                            let avail_files = std::fs::read_dir(sfx_dir.join("ambient")).unwrap();
                            let chosen_file = avail_files.choose(&mut rand::rng()).unwrap().unwrap();
                            player.play_sample(&chosen_file.path()).await;
                        }
                    }
                }
            }
        }
    });

    SfxControl {
        sink: event_sink
    }
}