zynko 0.1.0

//! A bot in a voice room. The Zynko server hands out a media token
//! (POST /channels/{id}/voice) and the server's members learn where the bot is from gateway
//! op 4; the audio itself travels through LiveKit, the same media server the apps use, through
//! the official `livekit` crate.
//!
//! ```no_run
//! # async fn f(client: zynko::Client, channel_id: &str) -> zynko::Result<()> {
//! use zynko::voice::JoinVoiceOptions;
//! let voice = client.join_voice(channel_id, JoinVoiceOptions::default()).await?;
//! voice.player().enqueue("https://exemplo.com/musica.mp3", serde_json::Value::Null);
//! # Ok(()) }
//! ```

use std::borrow::Cow;
use std::sync::{Arc, Mutex, Weak};
use std::time::Duration;

use async_trait::async_trait;
use futures_util::StreamExt;
use livekit::options::{AudioEncoding, TrackPublishOptions};
use livekit::prelude::*;
use livekit::webrtc::audio_frame::AudioFrame;
use livekit::webrtc::audio_source::native::NativeAudioSource;
use livekit::webrtc::audio_source::AudioSourceOptions;
use livekit::webrtc::audio_stream::native::NativeAudioStream;
use livekit::webrtc::peer_connection_factory::{ContinualGatheringPolicy, IceTransportsType};
use serde_json::json;
use tokio::sync::{broadcast, mpsc, watch};

use crate::client::Client;
use crate::error::{Error, Result};
use crate::player::{AudioPlayer, FrameSink, PlayerOptions, CHANNELS, FRAME_SAMPLES, SAMPLE_RATE};
use crate::types::{Channel, Id, VoiceState};

/// How to join a voice room.
#[derive(Debug, Clone, Default)]
pub struct JoinVoiceOptions {
    /// Hear the room: [`VoiceEvent::Audio`] with each person's sound (48 kHz mono). Off by
    /// default, so a bot that only plays does not download everyone's audio.
    pub receive: bool,
    /// "Conexão alternativa": media only through the voice server's relay.
    pub relay: bool,
    /// Show the bot as deafened in the room (it still plays).
    pub deaf: bool,
    /// ffmpeg and yt-dlp for the player.
    pub player: PlayerOptions,
}

/// Where a voice connection is.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VoiceStatus {
    Connecting,
    Connected,
    /// Moving to another voice server.
    Reconnecting,
    Closed,
}

/// What a voice connection tells.
#[derive(Debug, Clone)]
pub enum VoiceEvent {
    /// Connected to the room's media (again, after a server move).
    Ready,
    /// One person's audio (with `receive`): interleaved 16-bit samples at 48 kHz, in 10 ms
    /// pieces.
    Audio {
        user_id: Id,
        samples: Arc<[i16]>,
        channels: u32,
    },
    /// Who is speaking now, by user id.
    Speaking(Vec<Id>),
    /// Out of the room for good: left, removed, or the room went away.
    Disconnect(String),
    Error(String),
}

struct VState {
    status: VoiceStatus,
    server: Option<String>,
    can_speak: bool,
    room: Option<Arc<Room>>,
    source: Option<NativeAudioSource>,
    /// Which room connection is current; events from older ones are ignored.
    generation: u64,
    leaving: bool,
}

struct VoiceInner {
    client: Client,
    channel_id: Id,
    opts: JoinVoiceOptions,
    player: AudioPlayer,
    state: Mutex<VState>,
    moving: watch::Sender<bool>,
    events: broadcast::Sender<VoiceEvent>,
}

/// The bot in a voice room. Cheap to clone. Get one from
/// [`Client::join_voice`](crate::Client::join_voice).
#[derive(Clone)]
pub struct VoiceConnection {
    inner: Arc<VoiceInner>,
}

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

/// The player's way into the connection, without keeping it alive.
struct Sink(Weak<VoiceInner>);

#[async_trait]
impl FrameSink for Sink {
    async fn capture(&self, frame: &[i16]) {
        let Some(inner) = self.0.upgrade() else {
            return;
        };
        // Moving to another server: wait for the new connection, then go on where it was.
        let mut moving = inner.moving.subscribe();
        let _ = moving.wait_for(|m| !*m).await;
        let source = {
            let s = inner.state.lock().unwrap_or_else(|e| e.into_inner());
            if s.status == VoiceStatus::Closed {
                None
            } else {
                s.source.clone()
            }
        };
        drop(inner);
        match source {
            Some(src) => {
                let f = AudioFrame {
                    data: Cow::Borrowed(frame),
                    sample_rate: SAMPLE_RATE,
                    num_channels: CHANNELS,
                    samples_per_channel: FRAME_SAMPLES,
                };
                let _ = src.capture_frame(&f).await;
            }
            // Not allowed to speak here: keep time anyway.
            None => tokio::time::sleep(Duration::from_millis(20)).await,
        }
    }

    fn flush(&self) {
        if let Some(inner) = self.0.upgrade() {
            if let Some(src) = inner
                .state
                .lock()
                .unwrap_or_else(|e| e.into_inner())
                .source
                .as_ref()
            {
                src.clear_buffer();
            }
        }
    }
}

impl VoiceConnection {
    pub(crate) fn new(client: Client, channel_id: Id, opts: JoinVoiceOptions) -> Self {
        let (moving, _) = watch::channel(false);
        let (events, _) = broadcast::channel(4096);
        let inner = Arc::new_cyclic(|weak: &Weak<VoiceInner>| VoiceInner {
            player: AudioPlayer::new(Arc::new(Sink(weak.clone())), opts.player.clone()),
            client,
            channel_id,
            opts,
            state: Mutex::new(VState {
                status: VoiceStatus::Connecting,
                server: None,
                can_speak: true,
                room: None,
                source: None,
                generation: 0,
                leaving: false,
            }),
            moving,
            events,
        });
        Self { inner }
    }

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

    pub(crate) fn same(&self, other: &VoiceConnection) -> bool {
        Arc::ptr_eq(&self.inner, &other.inner)
    }

    /// The room.
    pub fn channel_id(&self) -> &str {
        &self.inner.channel_id
    }

    /// The room's channel, when the client has seen it.
    pub fn channel(&self) -> Option<Channel> {
        self.inner.client.channel(&self.inner.channel_id)
    }

    /// The player.
    pub fn player(&self) -> &AudioPlayer {
        &self.inner.player
    }

    pub fn status(&self) -> VoiceStatus {
        self.state().status
    }

    pub fn is_closed(&self) -> bool {
        self.status() == VoiceStatus::Closed
    }

    /// The voice server the room is on (Brasil 1, Brasil 2...).
    pub fn server(&self) -> Option<String> {
        self.state().server.clone()
    }

    /// Can speak here: a room's roles may let the bot only listen.
    pub fn can_speak(&self) -> bool {
        self.state().can_speak
    }

    /// Ready, audio, speaking, disconnect and error events.
    pub fn events(&self) -> broadcast::Receiver<VoiceEvent> {
        self.inner.events.subscribe()
    }

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

    /// Connects to the room's media and tells the server the bot is there.
    pub(crate) async fn connect(&self) -> Result<()> {
        let join = self
            .inner
            .client
            .rest()
            .voice_token(&self.inner.channel_id)
            .await?;
        let can_speak = join.can_speak != Some(false);
        {
            let mut s = self.state();
            s.server = join.server.clone();
            s.can_speak = can_speak;
        }
        let mut options = RoomOptions::default();
        options.auto_subscribe = self.inner.opts.receive;
        options.dynacast = false;
        if self.inner.opts.relay {
            options.rtc_config.ice_transport_type = IceTransportsType::Relay;
            options.rtc_config.continual_gathering_policy =
                ContinualGatheringPolicy::GatherContinually;
            options.rtc_config.ice_servers = Vec::new();
        }
        let (room, rx) = Room::connect(&join.url, &join.token, options)
            .await
            .map_err(|e| Error::Voice(format!("não conectou à mídia da sala: {e}")))?;
        let room = Arc::new(room);
        let generation = {
            let mut s = self.state();
            if s.status == VoiceStatus::Closed {
                None
            } else {
                s.generation += 1;
                s.room = Some(room.clone());
                Some(s.generation)
            }
        };
        let Some(generation) = generation else {
            // Left while connecting.
            let _ = room.close().await;
            return Ok(());
        };
        tokio::spawn(self.clone().room_events(generation, rx));
        if can_speak {
            let source =
                NativeAudioSource::new(AudioSourceOptions::default(), SAMPLE_RATE, CHANNELS, 200);
            let track = LocalAudioTrack::create_audio_track(
                "zynko",
                RtcAudioSource::Native(source.clone()),
            );
            let publish = TrackPublishOptions {
                source: TrackSource::Microphone,
                dtx: false,
                red: false,
                audio_encoding: Some(AudioEncoding {
                    max_bitrate: 128_000,
                }),
                ..Default::default()
            };
            room.local_participant()
                .publish_track(LocalTrack::Audio(track), publish)
                .await
                .map_err(|e| Error::Voice(format!("não publicou o áudio: {e}")))?;
            self.state().source = Some(source);
        }
        self.announce();
        {
            let mut s = self.state();
            if s.status == VoiceStatus::Closed {
                return Ok(());
            }
            s.status = VoiceStatus::Connected;
        }
        self.emit(VoiceEvent::Ready);
        Ok(())
    }

    async fn room_events(self, generation: u64, mut rx: mpsc::UnboundedReceiver<RoomEvent>) {
        while let Some(ev) = rx.recv().await {
            let current = {
                let s = self.state();
                s.generation == generation && s.status != VoiceStatus::Closed
            };
            if !current {
                return;
            }
            match ev {
                RoomEvent::Disconnected { reason } => {
                    let (leaving, moving) = (self.state().leaving, *self.inner.moving.borrow());
                    // The room closed under us: a server move is followed through CHANNEL_UPDATE;
                    // anything else (removed, room deleted) ends the connection.
                    if !leaving && !moving {
                        self.end(format!("desconectado ({reason:?})")).await;
                    }
                    return;
                }
                RoomEvent::ActiveSpeakersChanged { speakers } => {
                    self.emit(VoiceEvent::Speaking(
                        speakers
                            .iter()
                            .map(|p| p.identity().as_str().to_string())
                            .collect(),
                    ));
                }
                RoomEvent::TrackSubscribed {
                    track: RemoteTrack::Audio(track),
                    participant,
                    ..
                } if self.inner.opts.receive => {
                    let user_id = participant.identity().as_str().to_string();
                    tokio::spawn(self.clone().hear(generation, track, user_id));
                }
                _ => {}
            }
        }
    }

    async fn hear(self, generation: u64, track: RemoteAudioTrack, user_id: Id) {
        let mut stream = NativeAudioStream::new(track.rtc_track(), SAMPLE_RATE as i32, 1);
        while let Some(frame) = stream.next().await {
            {
                let s = self.state();
                if s.generation != generation || s.status == VoiceStatus::Closed {
                    break;
                }
            }
            if self.inner.events.receiver_count() == 0 {
                continue;
            }
            self.emit(VoiceEvent::Audio {
                user_id: user_id.clone(),
                samples: Arc::from(frame.data.as_ref()),
                channels: frame.num_channels,
            });
        }
    }

    /// Leaves the room.
    pub async fn leave(&self) {
        {
            let mut s = self.state();
            if s.status == VoiceStatus::Closed {
                return;
            }
            s.leaving = true;
        }
        self.inner.player.stop();
        self.inner.client.send_op(
            4,
            json!({ "channel_id": null, "self_mute": false, "self_deaf": false }),
        );
        self.end("saiu".into()).await;
    }

    // ---------- from the client ----------

    /// The room moved to another voice server (someone picked Brasil 1 or 2): follow it.
    pub(crate) fn on_channel_update(&self, c: &Channel) {
        {
            let s = self.state();
            let Some(new) = c.voice_server.as_ref() else {
                return;
            };
            if c.id != self.inner.channel_id
                || s.status != VoiceStatus::Connected
                || Some(new) == s.server.as_ref()
            {
                return;
            }
        }
        let me = self.clone();
        self.inner.moving.send_replace(true);
        tokio::spawn(async move {
            let (old, old_source) = {
                let mut s = me.state();
                s.status = VoiceStatus::Reconnecting;
                s.generation += 1;
                (s.room.take(), s.source.take())
            };
            if let Some(src) = old_source {
                src.clear_buffer();
            }
            if let Some(room) = old {
                let _ = room.close().await;
            }
            let result = Box::pin(me.connect()).await;
            me.inner.moving.send_replace(false);
            if let Err(e) = result {
                me.end("não conseguiu entrar no servidor novo".into()).await;
                me.emit(VoiceEvent::Error(e.to_string()));
            }
        });
    }

    /// The server says where the bot is in voice: elsewhere, or nowhere, ends this connection.
    pub(crate) fn on_own_state(&self, v: &VoiceState) {
        {
            let s = self.state();
            if s.status == VoiceStatus::Closed || s.leaving {
                return;
            }
        }
        if v.channel_id.as_deref() == Some(self.inner.channel_id.as_str()) {
            return;
        }
        // Leaving a seat an older gateway connection held (a restart within the grace period)
        // is not about this one.
        if v.channel_id.is_none()
            && !v.session_id.is_empty()
            && v.session_id != self.inner.client.cache().session_id
        {
            return;
        }
        let me = self.clone();
        let reason = if v.channel_id.is_some() {
            "entrou em outra sala"
        } else {
            "removido da sala"
        };
        tokio::spawn(async move { me.end(reason.into()).await });
    }

    /// After the gateway reconnects: say again where the bot is.
    pub(crate) fn resync(&self) {
        if matches!(
            self.status(),
            VoiceStatus::Connected | VoiceStatus::Reconnecting
        ) {
            self.announce();
        }
    }

    fn announce(&self) {
        let can_speak = self.state().can_speak;
        self.inner.client.send_op(
            4,
            json!({ "channel_id": self.inner.channel_id, "self_mute": !can_speak, "self_deaf": self.inner.opts.deaf }),
        );
    }

    async fn end(&self, reason: String) {
        let (room, source) = {
            let mut s = self.state();
            if s.status == VoiceStatus::Closed {
                return;
            }
            s.status = VoiceStatus::Closed;
            s.generation += 1;
            (s.room.take(), s.source.take())
        };
        self.inner.moving.send_replace(false);
        self.inner.player.stop();
        if let Some(src) = source {
            src.clear_buffer();
        }
        if let Some(room) = room {
            let _ = room.close().await;
        }
        self.inner.client.clear_voice(self);
        self.emit(VoiceEvent::Disconnect(reason));
    }
}