zynko 0.1.0

//! Plays audio into a voice connection: files and links through ffmpeg (any format and stream
//! it opens: mp3, ogg, wav, radio streams...), encoded streams, raw PCM, or tones. A queue,
//! pause, skip, volume. Audio goes out as 48 kHz stereo 16-bit, in 20 ms frames, paced by the
//! sink (the voice connection takes a frame only when there is room for it).

use std::collections::VecDeque;
use std::path::PathBuf;
use std::process::Stdio;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};

use async_trait::async_trait;
use serde_json::Value;
use tokio::io::{AsyncRead, AsyncReadExt};
use tokio::process::{Child, Command};
use tokio::sync::{broadcast, watch};

use crate::error::{Error, Result};

pub const SAMPLE_RATE: u32 = 48_000;
pub const CHANNELS: u32 = 2;
/// 20 ms per channel.
pub const FRAME_SAMPLES: u32 = 960;
const FRAME_LEN: usize = (FRAME_SAMPLES * CHANNELS) as usize;
const FRAME_BYTES: usize = FRAME_LEN * 2;

/// Something that reads bytes, for [`AudioInput::Stream`] and [`AudioInput::Pcm`].
pub type Reader = Box<dyn AsyncRead + Send + Unpin + 'static>;

/// What to play.
pub enum AudioInput {
    /// A path or an http(s) link, opened by ffmpeg.
    Url(String),
    /// Encoded audio (an mp3 download, a pipe), decoded by ffmpeg.
    Stream(Reader),
    /// Raw PCM already at 48 kHz, stereo, 16-bit little-endian.
    Pcm(Reader),
}

impl std::fmt::Debug for AudioInput {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            AudioInput::Url(u) => f.debug_tuple("Url").field(u).finish(),
            AudioInput::Stream(_) => f.write_str("Stream(..)"),
            AudioInput::Pcm(_) => f.write_str("Pcm(..)"),
        }
    }
}

impl From<&str> for AudioInput {
    fn from(s: &str) -> Self {
        AudioInput::Url(s.to_string())
    }
}
impl From<String> for AudioInput {
    fn from(s: String) -> Self {
        AudioInput::Url(s)
    }
}
impl From<PathBuf> for AudioInput {
    fn from(p: PathBuf) -> Self {
        AudioInput::Url(p.to_string_lossy().into_owned())
    }
}
impl From<&std::path::Path> for AudioInput {
    fn from(p: &std::path::Path) -> Self {
        AudioInput::Url(p.to_string_lossy().into_owned())
    }
}

/// A sine wave, for testing: `player.play(tone(440.0, 2.0), Value::Null)`.
pub fn tone(hz: f64, seconds: f64) -> AudioInput {
    tone_with_volume(hz, seconds, 0.2)
}

/// A sine wave at a given loudness (0 to 1).
pub fn tone_with_volume(hz: f64, seconds: f64, volume: f64) -> AudioInput {
    AudioInput::Pcm(Box::new(std::io::Cursor::new(tone_pcm(
        hz, seconds, volume,
    ))))
}

pub(crate) fn tone_pcm(hz: f64, seconds: f64, volume: f64) -> Vec<u8> {
    let total = (seconds * SAMPLE_RATE as f64).round().max(0.0) as usize;
    let mut out = Vec::with_capacity(total * 4);
    for i in 0..total {
        let v = ((2.0 * std::f64::consts::PI * hz * i as f64 / SAMPLE_RATE as f64).sin()
            * volume
            * 32767.0)
            .round() as i16;
        let b = v.to_le_bytes();
        out.extend_from_slice(&b);
        out.extend_from_slice(&b);
    }
    out
}

/// Waits until a flag has a value.
async fn until(rx: &mut watch::Receiver<bool>, value: bool) {
    let _ = rx.wait_for(|v| *v == value).await;
}

/// Scales samples by `v`, clipping.
pub(crate) fn scale(frame: &mut [i16], v: f32) {
    for s in frame.iter_mut() {
        *s = (*s as f32 * v).round().clamp(-32768.0, 32767.0) as i16;
    }
}

/// A queued or playing track. Cheap to clone; `meta` is whatever you want back in events (a
/// title, who asked for it...).
#[derive(Debug, Clone)]
pub struct Track {
    /// Unique per player.
    pub id: u64,
    pub meta: Value,
    /// The link or path, for `Url` inputs.
    pub source: Option<String>,
}

/// What the player tells.
#[derive(Debug, Clone)]
pub enum PlayerEvent {
    Start(Track),
    Finish(Track),
    Error {
        error: String,
        track: Track,
    },
    /// Nothing left in the queue.
    Idle,
}

/// Player settings.
#[derive(Debug, Clone, Default)]
pub struct PlayerOptions {
    /// ffmpeg's path. Default: FFMPEG_PATH, or `ffmpeg` on the PATH.
    pub ffmpeg: Option<String>,
    /// Hand http(s) links to yt-dlp first (its path, or `Some("yt-dlp")`), for pages of sites
    /// it supports instead of direct audio links.
    pub ytdlp: Option<String>,
}

/// Where the player writes frames: the voice connection.
#[async_trait]
pub trait FrameSink: Send + Sync + 'static {
    /// Takes one 20 ms frame (interleaved stereo); returns when there is room for the next.
    async fn capture(&self, frame: &[i16]);
    /// Drops what is buffered, so a skip or pause is heard at once.
    fn flush(&self);
}

struct Entry {
    track: Track,
    input: AudioInput,
}

struct State {
    queue: VecDeque<Entry>,
    current: Option<Track>,
    running: bool,
    volume: f32,
    /// Cancels the current track.
    abort: Option<watch::Sender<bool>>,
}

struct PlayerInner {
    sink: Arc<dyn FrameSink>,
    opts: PlayerOptions,
    state: Mutex<State>,
    paused: watch::Sender<bool>,
    events: broadcast::Sender<PlayerEvent>,
    next_id: AtomicU64,
}

/// The queue and the playback. Cheap to clone.
///
/// ```no_run
/// # fn f(p: zynko::AudioPlayer) {
/// use serde_json::json;
/// p.enqueue("https://radio.exemplo.com/stream", json!({ "title": "Rádio" }));
/// p.play(zynko::tone(440.0, 3.0), json!(null));
/// p.set_volume(0.5);
/// p.skip();
/// # }
/// ```
#[derive(Clone)]
pub struct AudioPlayer {
    inner: Arc<PlayerInner>,
}

impl std::fmt::Debug for AudioPlayer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AudioPlayer")
            .field("current", &self.current())
            .finish_non_exhaustive()
    }
}

impl AudioPlayer {
    /// A player writing into `sink`. A voice connection makes its own; this is for tests and
    /// other sinks.
    pub fn new(sink: Arc<dyn FrameSink>, opts: PlayerOptions) -> Self {
        let (paused, _) = watch::channel(false);
        let (events, _) = broadcast::channel(256);
        Self {
            inner: Arc::new(PlayerInner {
                sink,
                opts,
                state: Mutex::new(State {
                    queue: VecDeque::new(),
                    current: None,
                    running: false,
                    volume: 1.0,
                    abort: None,
                }),
                paused,
                events,
                next_id: AtomicU64::new(1),
            }),
        }
    }

    fn state(&self) -> std::sync::MutexGuard<'_, State> {
        self.inner.state.lock().unwrap_or_else(|e| e.into_inner())
    }

    /// Start, finish, error and idle events.
    pub fn events(&self) -> broadcast::Receiver<PlayerEvent> {
        self.inner.events.subscribe()
    }

    fn emit(&self, ev: PlayerEvent) {
        let _ = self.inner.events.send(ev);
    }

    fn entry(&self, input: AudioInput, meta: Value) -> Entry {
        let source = match &input {
            AudioInput::Url(u) => Some(u.clone()),
            _ => None,
        };
        let track = Track {
            id: self.inner.next_id.fetch_add(1, Ordering::Relaxed),
            meta,
            source,
        };
        Entry { track, input }
    }

    /// Plays now, in place of what was playing (the queue stays).
    pub fn play(&self, input: impl Into<AudioInput>, meta: Value) -> Track {
        let e = self.entry(input.into(), meta);
        let track = e.track.clone();
        let start = {
            let mut s = self.state();
            s.queue.push_front(e);
            !std::mem::replace(&mut s.running, true)
        };
        if start {
            self.spawn_loop();
        } else {
            self.skip();
        }
        track
    }

    /// Adds to the end of the queue; starts if nothing is playing.
    pub fn enqueue(&self, input: impl Into<AudioInput>, meta: Value) -> Track {
        let e = self.entry(input.into(), meta);
        let track = e.track.clone();
        let start = {
            let mut s = self.state();
            s.queue.push_back(e);
            !std::mem::replace(&mut s.running, true)
        };
        if start {
            self.spawn_loop();
        }
        track
    }

    /// Goes to the next in the queue.
    pub fn skip(&self) {
        if let Some(a) = self.state().abort.as_ref() {
            let _ = a.send(true);
        }
        self.inner.sink.flush();
        self.resume();
    }

    /// Stops and empties the queue.
    pub fn stop(&self) {
        self.state().queue.clear();
        self.skip();
    }

    pub fn pause(&self) {
        if *self.inner.paused.borrow() || self.state().current.is_none() {
            return;
        }
        self.inner.paused.send_replace(true);
        self.inner.sink.flush();
    }

    pub fn resume(&self) {
        self.inner.paused.send_replace(false);
    }

    pub fn is_paused(&self) -> bool {
        *self.inner.paused.borrow()
    }

    /// 0 (silent) to 2 (twice as loud); 1 is as it came.
    pub fn volume(&self) -> f32 {
        self.state().volume
    }

    pub fn set_volume(&self, v: f32) {
        self.state().volume = if v.is_nan() { 1.0 } else { v.clamp(0.0, 2.0) };
    }

    /// What is playing now.
    pub fn current(&self) -> Option<Track> {
        self.state().current.clone()
    }

    /// Played after the current one, in order.
    pub fn queue(&self) -> Vec<Track> {
        self.state().queue.iter().map(|e| e.track.clone()).collect()
    }

    fn spawn_loop(&self) {
        let me = self.clone();
        tokio::spawn(async move { me.run().await });
    }

    async fn run(self) {
        loop {
            let (entry, abort_rx) = {
                let mut s = self.state();
                match s.queue.pop_front() {
                    Some(e) => {
                        let (tx, rx) = watch::channel(false);
                        s.abort = Some(tx);
                        s.current = Some(e.track.clone());
                        (e, rx)
                    }
                    None => {
                        s.current = None;
                        s.running = false;
                        s.abort = None;
                        drop(s);
                        self.inner.paused.send_replace(false);
                        self.emit(PlayerEvent::Idle);
                        return;
                    }
                }
            };
            let track = entry.track.clone();
            self.emit(PlayerEvent::Start(track.clone()));
            match self.play_one(entry.input, abort_rx).await {
                Ok(()) => self.emit(PlayerEvent::Finish(track)),
                Err(e) => self.emit(PlayerEvent::Error {
                    error: e.to_string(),
                    track,
                }),
            }
        }
    }

    async fn play_one(&self, input: AudioInput, mut abort: watch::Receiver<bool>) -> Result<()> {
        let (mut reader, mut children) = self.open(input)?;
        let mut buf = vec![0u8; FRAME_BYTES];
        let mut filled = 0usize;
        let mut frame = vec![0i16; FRAME_LEN];
        let mut stopped = false;
        let mut paused = self.inner.paused.subscribe();
        'read: loop {
            let r = tokio::select! {
                biased;
                _ = until(&mut abort, true) => { stopped = true; break 'read; }
                r = reader.read(&mut buf[filled..]) => r,
            };
            let n = match r {
                Ok(n) => n,
                // A broken pipe after a kill, or a dead source: ends the track.
                Err(e) => {
                    stopped = *abort.borrow();
                    if stopped {
                        break 'read;
                    }
                    return Err(e.into());
                }
            };
            if n == 0 {
                break;
            }
            filled += n;
            if filled < FRAME_BYTES {
                continue;
            }
            filled = 0;
            tokio::select! {
                biased;
                _ = until(&mut abort, true) => { stopped = true; break 'read; }
                _ = until(&mut paused, false) => {}
            }
            for (i, s) in frame.iter_mut().enumerate() {
                *s = i16::from_le_bytes([buf[i * 2], buf[i * 2 + 1]]);
            }
            let volume = self.state().volume;
            if volume != 1.0 {
                scale(&mut frame, volume);
            }
            tokio::select! {
                biased;
                _ = until(&mut abort, true) => { stopped = true; break 'read; }
                _ = self.inner.sink.capture(&frame) => {}
            }
        }
        // The last partial frame, padded with silence.
        if !stopped && filled > 0 {
            let whole = filled / 2;
            for (i, s) in frame.iter_mut().enumerate() {
                *s = if i < whole {
                    i16::from_le_bytes([buf[i * 2], buf[i * 2 + 1]])
                } else {
                    0
                };
            }
            self.inner.sink.capture(&frame).await;
        }
        // Skipped: kill first, so ffmpeg does not complain about the pipe closing under it.
        if stopped {
            for (_, child) in children.iter_mut() {
                let _ = child.start_kill();
            }
            for (_, child) in children.iter_mut() {
                let _ = child.wait().await;
            }
            return Ok(());
        }
        drop(reader);
        let mut failed = None;
        for (name, child) in children.iter_mut() {
            match child.wait().await {
                Ok(st) if !st.success() => {
                    failed = Some(format!(
                        "{name} terminou com código {}",
                        st.code().unwrap_or(-1)
                    ))
                }
                _ => {}
            }
        }
        match failed {
            Some(msg) => Err(Error::Voice(msg)),
            None => Ok(()),
        }
    }

    /// The input as raw PCM, and the processes making it.
    fn open(&self, input: AudioInput) -> Result<(Reader, Vec<(String, Child)>)> {
        let source: Option<Reader>;
        let mut children = Vec::new();
        let ffmpeg = self
            .inner
            .opts
            .ffmpeg
            .clone()
            .or_else(|| std::env::var("FFMPEG_PATH").ok())
            .unwrap_or_else(|| "ffmpeg".into());
        let out = [
            "-f",
            "s16le",
            "-ar",
            "48000",
            "-ac",
            "2",
            "-loglevel",
            "error",
            "pipe:1",
        ];
        let mut args: Vec<String> = vec!["-hide_banner".into()];
        match input {
            AudioInput::Pcm(r) => return Ok((r, children)),
            AudioInput::Url(url) => {
                let http = url.len() > 8
                    && (url[..7].eq_ignore_ascii_case("http://")
                        || url[..8].eq_ignore_ascii_case("https://"));
                if let (true, Some(ytdlp)) = (http, self.inner.opts.ytdlp.as_ref()) {
                    let mut y = Command::new(ytdlp)
                        .args([
                            "-f",
                            "bestaudio/best",
                            "--no-playlist",
                            "--quiet",
                            "-o",
                            "-",
                            &url,
                        ])
                        .stdin(Stdio::null())
                        .stdout(Stdio::piped())
                        .stderr(Stdio::inherit())
                        .kill_on_drop(true)
                        .spawn()
                        .map_err(|e| {
                            Error::Voice(format!("não deu para abrir o yt-dlp ({ytdlp}): {e}"))
                        })?;
                    source = y.stdout.take().map(|s| Box::new(s) as Reader);
                    children.push(("yt-dlp".to_string(), y));
                    args.extend(["-i".into(), "pipe:0".into()]);
                } else {
                    source = None;
                    args.push("-nostdin".into());
                    if http {
                        args.extend(
                            [
                                "-reconnect",
                                "1",
                                "-reconnect_streamed",
                                "1",
                                "-reconnect_delay_max",
                                "5",
                            ]
                            .map(String::from),
                        );
                    }
                    args.extend(["-i".into(), url]);
                }
            }
            AudioInput::Stream(r) => {
                source = Some(r);
                args.extend(["-i".into(), "pipe:0".into()]);
            }
        }
        args.extend(out.map(String::from));
        let mut f = Command::new(&ffmpeg)
            .args(&args)
            .stdin(if source.is_some() {
                Stdio::piped()
            } else {
                Stdio::null()
            })
            .stdout(Stdio::piped())
            .stderr(Stdio::inherit())
            .kill_on_drop(true)
            .spawn()
            .map_err(|e| Error::Voice(format!("não deu para abrir o ffmpeg ({ffmpeg}): {e}")))?;
        if let (Some(mut src), Some(mut stdin)) = (source, f.stdin.take()) {
            tokio::spawn(async move {
                // Ends when the source ends, or when ffmpeg goes away (skip).
                let _ = tokio::io::copy(&mut src, &mut stdin).await;
            });
        }
        let stdout = f
            .stdout
            .take()
            .ok_or_else(|| Error::Voice("ffmpeg sem saída".into()))?;
        children.push(("ffmpeg".to_string(), f));
        Ok((Box::new(stdout), children))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::AtomicUsize;
    use std::time::Duration;

    #[test]
    fn tone_is_stereo_s16le() {
        let pcm = tone_pcm(440.0, 0.5, 0.2);
        assert_eq!(pcm.len(), 24_000 * 4);
        let s = |i: usize| i16::from_le_bytes([pcm[i * 2], pcm[i * 2 + 1]]);
        assert_eq!(s(0), 0);
        // Left and right are the same; peaks reach about 0.2 of full scale.
        assert_eq!(s(10), s(11));
        let peak = (0..pcm.len() / 2)
            .map(s)
            .map(|v| v.unsigned_abs())
            .max()
            .unwrap();
        assert!((6500..=6554).contains(&peak), "{peak}");
    }

    #[test]
    fn scale_clips() {
        let mut f = [1000, -1000, 30000, -30000];
        scale(&mut f, 2.0);
        assert_eq!(f, [2000, -2000, 32767, -32768]);
        scale(&mut f, 0.0);
        assert_eq!(f, [0, 0, 0, 0]);
    }

    /// Counts frames and paces them like a real sink would (but faster).
    struct Counter {
        frames: AtomicUsize,
        loud: AtomicUsize,
    }

    #[async_trait]
    impl FrameSink for Counter {
        async fn capture(&self, frame: &[i16]) {
            assert_eq!(frame.len(), FRAME_LEN);
            self.frames.fetch_add(1, Ordering::Relaxed);
            if frame.iter().any(|&s| s.unsigned_abs() > 10_000) {
                self.loud.fetch_add(1, Ordering::Relaxed);
            }
            tokio::time::sleep(Duration::from_millis(1)).await;
        }
        fn flush(&self) {}
    }

    #[tokio::test]
    async fn plays_the_queue_in_order_with_volume() {
        let sink = Arc::new(Counter {
            frames: AtomicUsize::new(0),
            loud: AtomicUsize::new(0),
        });
        let p = AudioPlayer::new(sink.clone(), PlayerOptions::default());
        let mut ev = p.events();
        p.set_volume(3.0);
        assert_eq!(p.volume(), 2.0);
        p.set_volume(1.0);
        let a = p.enqueue(tone_with_volume(440.0, 0.1, 0.5), Value::from("a"));
        let b = p.enqueue(tone_with_volume(440.0, 0.105, 0.5), Value::from("b"));
        let mut order = Vec::new();
        loop {
            match tokio::time::timeout(Duration::from_secs(10), ev.recv())
                .await
                .unwrap()
                .unwrap()
            {
                PlayerEvent::Start(t) => order.push(format!("start {}", t.id)),
                PlayerEvent::Finish(t) => order.push(format!("finish {}", t.id)),
                PlayerEvent::Error { error, .. } => panic!("{error}"),
                PlayerEvent::Idle => break,
            }
        }
        assert_eq!(
            order,
            [
                format!("start {}", a.id),
                format!("finish {}", a.id),
                format!("start {}", b.id),
                format!("finish {}", b.id)
            ]
        );
        // 0.1 s = 5 frames; 0.105 s = 5 frames and a padded sixth.
        assert_eq!(sink.frames.load(Ordering::Relaxed), 11);
        // At volume 0.5 of full scale, peaks pass 10 000.
        assert!(sink.loud.load(Ordering::Relaxed) > 0);
        assert!(p.current().is_none());
    }

    #[tokio::test]
    async fn skip_stop_pause() {
        let sink = Arc::new(Counter {
            frames: AtomicUsize::new(0),
            loud: AtomicUsize::new(0),
        });
        let p = AudioPlayer::new(sink.clone(), PlayerOptions::default());
        let mut ev = p.events();
        p.enqueue(tone(440.0, 30.0), Value::Null);
        p.enqueue(tone(440.0, 30.0), Value::Null);
        p.enqueue(tone(440.0, 30.0), Value::Null);
        assert!(matches!(ev.recv().await.unwrap(), PlayerEvent::Start(_)));
        assert_eq!(p.queue().len(), 2);
        p.pause();
        assert!(p.is_paused());
        tokio::time::sleep(Duration::from_millis(50)).await;
        let n = sink.frames.load(Ordering::Relaxed);
        tokio::time::sleep(Duration::from_millis(50)).await;
        assert!(
            sink.frames.load(Ordering::Relaxed) <= n + 1,
            "paused but still playing"
        );
        p.resume();
        p.skip();
        assert!(matches!(ev.recv().await.unwrap(), PlayerEvent::Finish(_)));
        assert!(matches!(ev.recv().await.unwrap(), PlayerEvent::Start(_)));
        assert_eq!(p.queue().len(), 1);
        p.stop();
        assert!(matches!(ev.recv().await.unwrap(), PlayerEvent::Finish(_)));
        assert!(matches!(ev.recv().await.unwrap(), PlayerEvent::Idle));
        assert!(p.queue().is_empty());
    }

    #[tokio::test]
    async fn play_jumps_the_queue() {
        let sink = Arc::new(Counter {
            frames: AtomicUsize::new(0),
            loud: AtomicUsize::new(0),
        });
        let p = AudioPlayer::new(sink, PlayerOptions::default());
        let mut ev = p.events();
        let first = p.enqueue(tone(440.0, 30.0), Value::Null);
        let later = p.enqueue(tone(440.0, 0.04), Value::Null);
        assert!(matches!(ev.recv().await.unwrap(), PlayerEvent::Start(t) if t.id == first.id));
        let now = p.play(tone(880.0, 0.04), Value::Null);
        assert!(matches!(ev.recv().await.unwrap(), PlayerEvent::Finish(t) if t.id == first.id));
        assert!(matches!(ev.recv().await.unwrap(), PlayerEvent::Start(t) if t.id == now.id));
        assert!(matches!(ev.recv().await.unwrap(), PlayerEvent::Finish(t) if t.id == now.id));
        assert!(matches!(ev.recv().await.unwrap(), PlayerEvent::Start(t) if t.id == later.id));
    }

    #[tokio::test]
    async fn missing_ffmpeg_is_an_error_event() {
        let sink = Arc::new(Counter {
            frames: AtomicUsize::new(0),
            loud: AtomicUsize::new(0),
        });
        let p = AudioPlayer::new(
            sink,
            PlayerOptions {
                ffmpeg: Some("ffmpeg-que-nao-existe".into()),
                ytdlp: None,
            },
        );
        let mut ev = p.events();
        p.enqueue("musica.mp3", Value::Null);
        assert!(matches!(ev.recv().await.unwrap(), PlayerEvent::Start(_)));
        assert!(matches!(
            ev.recv().await.unwrap(),
            PlayerEvent::Error { .. }
        ));
        assert!(matches!(ev.recv().await.unwrap(), PlayerEvent::Idle));
    }
}