dj-kitty-cat/src/queue.rs

//! Implements a queue for the bot to play songs in order

use std::{collections::VecDeque, sync::Arc};

use anyhow::Result;
use parking_lot::Mutex;
use poise::async_trait;
use songbird::{
    events::EventData,
    input::Input,
    tracks::{self, Track, TrackHandle},
    Driver, Event, EventContext, EventHandler, TrackEvent,
};
use thiserror::Error;

#[derive(Error, Debug)]
pub enum QueueError {
    #[error("Nothing is in the queue.")]
    EmptyQueue,
}

/// Inner queue data behind a Mutex to allow the event handler to access it
struct QueueCore {
    tracks: VecDeque<TrackHandle>,
}

/// Event handler for the queue
struct QueueHandler {
    remote_lock: Arc<Mutex<QueueCore>>,
}

impl QueueHandler {
    /// Event to remove the track from the queue when it ends
    pub fn register_track_end_event(track: &mut Track, remote_lock: Arc<Mutex<QueueCore>>) {
        let position = track.position();
        track
            .events
            .as_mut()
            .expect("Why is this even an Option?")
            .add_event(
                EventData::new(Event::Track(TrackEvent::End), QueueHandler { remote_lock }),
                position,
            );
    }
}

#[async_trait]
impl EventHandler for QueueHandler {
    async fn act(&self, ctx: &EventContext<'_>) -> Option<Event> {
        let mut inner = self.remote_lock.lock();

        // Due to possibility that users might remove, reorder,
        // or dequeue+stop tracks, we need to verify that the FIRST
        // track is the one who has ended.
        match ctx {
            EventContext::Track(ts) => {
                // This slice should have exactly one entry.
                // If the ended track has same id as the queue head, then
                // we can progress the queue.
                if inner.tracks.front()?.uuid() != ts.first()?.1.uuid() {
                    return None;
                }
            }
            _ => return None,
        }

        let _old = inner.tracks.pop_front();

        // Keep going until we find one track which works, or we run out.
        while let Some(new) = inner.tracks.front() {
            if new.play().is_err() {
                // Discard files which cannot be used for whatever reason.
                inner.tracks.pop_front();
            } else {
                break;
            }
        }

        None
    }
}

/// A queue of tracks to play
pub struct Queue {
    inner: Arc<Mutex<QueueCore>>,
}

impl Queue {
    #[must_use]
    pub fn new() -> Self {
        Self {
            inner: Arc::new(Mutex::new(QueueCore {
                tracks: VecDeque::new(),
            })),
        }
    }

    /// Resumes the current track
    pub fn resume(&mut self) -> Result<()> {
        let inner = self.inner.lock();
        let Some(track) = inner.tracks.front() else {
            return Err(QueueError::EmptyQueue.into());
        };
        track.play()?;
        Ok(())
    }

    /// Stops the current track and removes it from the queue
    pub fn stop(&mut self) -> Result<()> {
        let mut inner = self.inner.lock();
        let Some(track) = inner.tracks.front() else {
            return Err(QueueError::EmptyQueue.into());
        };
        track.stop()?;
        inner.tracks.pop_front();
        Ok(())
    }

    /// Pauses the current track
    pub fn pause(&mut self) -> Result<()> {
        let inner = self.inner.lock();
        let Some(track) = inner.tracks.front() else {
            return Err(QueueError::EmptyQueue.into());
        };
        track.pause()?;
        Ok(())
    }

    /// Adds a track to the end of the queue
    pub fn add_to_end(&mut self, source: Input, handler: &mut Driver) -> TrackHandle {
        let mut inner = self.inner.lock();
        let (mut track, handle) = tracks::create_player(source);
        track.pause();
        QueueHandler::register_track_end_event(&mut track, self.inner.clone());
        inner.tracks.push_back(handle.clone());
        handler.play(track);
        handle
    }

    /// Adds multiple tracks to the end of the queue
    pub fn add_all_to_end(
        &mut self,
        sources: Vec<Input>,
        handler: &mut Driver,
    ) -> Vec<TrackHandle> {
        let mut handles = Vec::new();
        for source in sources {
            handles.push(self.add_to_end(source, handler));
        }
        handles
    }

    /// Adds a track to play next
    pub fn add_next(&mut self, source: Input, handler: &mut Driver) -> TrackHandle {
        let mut inner = self.inner.lock();
        let (mut track, handle) = tracks::create_player(source);
        track.pause();
        QueueHandler::register_track_end_event(&mut track, self.inner.clone());
        if inner.tracks.is_empty() {
            inner.tracks.push_back(handle.clone());
        } else {
            inner.tracks.insert(1, handle.clone());
        }
        handler.play(track);
        handle
    }

    /// Clears the queue
    pub fn clear(&mut self) {
        let mut inner = self.inner.lock();
        for track in inner.tracks.drain(..) {
            let _ = track.stop();
        }
    }

    /// Returns whether the queue is empty
    #[must_use]
    pub fn is_empty(&self) -> bool {
        let inner = self.inner.lock();
        inner.tracks.is_empty()
    }
}