path: root/crates/asist-audio/src/playback.rs

//! Playback of the synthesized audio.
//!
//! The TTS produces bursts (one codec block at a time) and the speaker consumes
//! at a constant rate, so there is a ring between them. The callback only
//! drains the ring; the synthesizer only fills it. Cutting the answer is then
//! a trivial operation with no race conditions: the ring is emptied and the
//! voice stops at the next audio block.

use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Condvar, Mutex};
use std::time::Duration;

use cpal::traits::{DeviceTrait, HostTrait, StreamTrait};
use cpal::SampleFormat;

use asist_core::config::AudioConfig;
use asist_core::error::{Error, Result};

use crate::{describe, to_mono_at, InputFormat, TTS_SAMPLE_RATE};

/// State shared between the synthesizer and the speaker callback.
struct Ring {
    samples: Mutex<std::collections::VecDeque<f32>>,
    /// Wakes whoever is waiting for the queue to empty.
    drained: Condvar,
    /// Samples served to the speaker since it started. Used to know whether
    /// anything has really played yet, which is the latency the user perceives.
    played: AtomicU64,
    /// There is audio in progress (queue left or still being fed).
    active: AtomicBool,
    gain: Mutex<f32>,
}

impl Ring {
    fn new(gain: f32) -> Self {
        Self {
            samples: Mutex::new(std::collections::VecDeque::new()),
            drained: Condvar::new(),
            played: AtomicU64::new(0),
            active: AtomicBool::new(false),
            gain: Mutex::new(gain),
        }
    }

    fn lock(&self) -> std::sync::MutexGuard<'_, std::collections::VecDeque<f32>> {
        self.samples.lock().unwrap_or_else(|e| e.into_inner())
    }
}

/// Speaker remote control: it can be cloned and shared across threads.
#[derive(Clone)]
pub struct PlaybackHandle {
    ring: Arc<Ring>,
    sample_rate: u32,
}

impl PlaybackHandle {
    /// Queues mono samples already at the device rate.
    pub fn push(&self, samples: &[f32]) {
        if samples.is_empty() {
            return;
        }
        let mut queue = self.ring.lock();
        queue.extend(samples.iter().copied());
        self.ring.active.store(true, Ordering::SeqCst);
    }

    /// Queues TTS audio (24 kHz mono), resampling if needed.
    pub fn push_tts(&self, samples: &[f32]) {
        if self.sample_rate == TTS_SAMPLE_RATE {
            self.push(samples);
            return;
        }
        let format = InputFormat {
            sample_rate: TTS_SAMPLE_RATE as usize,
            channels: 1,
        };
        self.push(&to_mono_at(samples, format, self.sample_rate));
    }

    /// Marks that feeding audio has finished and everything has played.
    ///
    /// It needs a method separate from `stop`: this one drops nothing, it only
    /// clears the «audio in progress» flag. Without it the flag stayed set after
    /// the first answer, the segmenter kept the microphone closed believing the
    /// assistant was still talking, and the assistant never heard anything again
    /// for the whole session.
    pub fn mark_idle(&self) {
        let queue = self.ring.lock();
        if queue.is_empty() {
            self.ring.active.store(false, Ordering::SeqCst);
        }
    }

    /// Goes silent right now and drops whatever was left to play.
    pub fn stop(&self) {
        let mut queue = self.ring.lock();
        queue.clear();
        self.ring.active.store(false, Ordering::SeqCst);
        self.ring.drained.notify_all();
    }

    /// Seconds of audio waiting to play.
    pub fn queued_secs(&self) -> f32 {
        self.ring.lock().len() as f32 / self.sample_rate as f32
    }

    /// `true` if audio has already come out of the speaker.
    pub fn has_played(&self) -> bool {
        self.ring.played.load(Ordering::SeqCst) > 0
    }

    pub fn reset_played(&self) {
        self.ring.played.store(0, Ordering::SeqCst);
    }

    pub fn is_active(&self) -> bool {
        self.ring.active.load(Ordering::SeqCst)
    }

    pub fn set_gain(&self, gain: f32) {
        *self.ring.gain.lock().unwrap_or_else(|e| e.into_inner()) = gain;
    }

    /// Waits for the queue to empty, or for the deadline.
    ///
    /// Returns `true` if everything finished playing and `false` if time ran
    /// out, so the caller can tell «done» from «still playing».
    pub fn wait_drained(&self, timeout: Duration) -> bool {
        let deadline = std::time::Instant::now() + timeout;
        let mut queue = self.ring.lock();
        while !queue.is_empty() {
            let left = deadline.saturating_duration_since(std::time::Instant::now());
            if left.is_zero() {
                return false;
            }
            let (guard, result) = self
                .ring
                .drained
                .wait_timeout(queue, left.min(Duration::from_millis(50)))
                .unwrap_or_else(|e| e.into_inner());
            queue = guard;
            if result.timed_out() && queue.is_empty() {
                break;
            }
        }
        self.ring.active.store(false, Ordering::SeqCst);
        true
    }
}

/// A handle with no device behind it, to test the ring logic without
/// opening a sound card.
pub fn detached_handle(sample_rate: u32) -> PlaybackHandle {
    PlaybackHandle {
        ring: Arc::new(Ring::new(1.0)),
        sample_rate,
    }
}

pub struct Playback {
    stream: cpal::Stream,
    handle: PlaybackHandle,
    pub device_name: String,
    pub sample_rate: u32,
}

impl Playback {
    pub fn open(config: &AudioConfig) -> Result<Self> {
        let host = cpal::default_host();
        let device = select_device(&host, &config.output_device)?;
        let device_name = describe(&device);

        let supported = preferred_config(&device)?;
        let sample_rate = supported.sample_rate();
        let channels = supported.channels() as usize;
        let stream_config: cpal::StreamConfig = supported.clone().into();

        let ring = Arc::new(Ring::new(config.output_gain));
        let on_error = |err| tracing::error!(target: "audio", %err, "output stream");

        // The callback: take whatever there is, fill the rest with silence
        // and leave. It never waits for audio, because waiting here is heard
        // as a dropout.
        macro_rules! build {
            ($sample:ty, $silence:expr, $convert:expr) => {{
                let ring = Arc::clone(&ring);
                device
                    .build_output_stream(
                        &stream_config,
                        move |data: &mut [$sample], _: &_| {
                            let gain = *ring.gain.lock().unwrap_or_else(|e| e.into_inner());
                            let mut queue = ring.lock();
                            let mut served = 0u64;
                            for frame in data.chunks_mut(channels) {
                                match queue.pop_front() {
                                    Some(sample) => {
                                        let value = (sample * gain).clamp(-1.0, 1.0);
                                        // Mono to N channels: the same sample on all of them.
                                        for out in frame.iter_mut() {
                                            *out = $convert(value);
                                        }
                                        served += 1;
                                    }
                                    None => {
                                        for out in frame.iter_mut() {
                                            *out = $silence;
                                        }
                                    }
                                }
                            }
                            if served > 0 {
                                ring.played.fetch_add(served, Ordering::SeqCst);
                            }
                            if queue.is_empty() {
                                ring.drained.notify_all();
                            }
                        },
                        on_error,
                        None,
                    )
                    .map_err(|e| Error::Audio(format!("could not open the output: {e}")))?
            }};
        }

        let stream = match supported.sample_format() {
            SampleFormat::F32 => build!(f32, 0.0f32, |v: f32| v),
            SampleFormat::I16 => build!(i16, 0i16, |v: f32| (v * 32767.0) as i16),
            SampleFormat::U16 => {
                build!(u16, u16::MAX / 2, |v: f32| ((v * 32767.0) as i32 + 32768)
                    as u16)
            }
            other => {
                return Err(Error::Audio(format!(
                    "unsupported output format: {other:?}"
                )))
            }
        };
        stream