FrenBot/src/discord/voices.rs

use crate::discord::Context;
use crate::error::Error;
use poise::async_trait;
use poise::serenity_prelude::all::CreateAttachment;
use poise::serenity_prelude::builder::CreateMessage;
use poise::serenity_prelude::model::id::UserId;
use poise::serenity_prelude::model::prelude::GuildId;
use poise::serenity_prelude::utils::MessageBuilder;
use songbird::driver::Bitrate;
use songbird::input::cached::Compressed;
use songbird::{Event, EventContext, EventHandler, Songbird, TrackEvent, input};
use std::borrow::Cow;
use std::collections::HashMap;
use std::fmt::{Display, Formatter};
use std::path::{Path, PathBuf};
use std::sync::Arc;

async fn get_voice_dictionary(path: &Path) -> Result<HashMap<String, PathBuf>, tokio::io::Error> {
    let mut dir = tokio::fs::read_dir(path).await?;
    let mut dict = HashMap::new();

    while let Some(word) = &dir.next_entry().await? {
        let word_name = word
            .path()
            .file_stem()
            .unwrap()
            .to_str()
            .unwrap()
            .to_string();
        dict.insert(word_name, word.path());
    }

    Ok(dict)
}

async fn get_voices(voice_path: &Path) -> Result<Vec<String>, tokio::io::Error> {
    let mut dir = tokio::fs::read_dir(voice_path).await?;
    let mut voices = Vec::new();

    while let Some(file) = &dir.next_entry().await? {
        if file.path().is_dir() {
            let voice_name = file
                .path()
                .file_name()
                .unwrap()
                .to_str()
                .unwrap()
                .to_string();

            voices.push(voice_name)
        }
    }

    Ok(voices)
}

async fn find_voice(voice_path: &Path, name: &str) -> Result<Option<PathBuf>, tokio::io::Error> {
    let mut dir = tokio::fs::read_dir(voice_path).await?;

    while let Some(file) = &dir.next_entry().await? {
        if file.path().is_dir() {
            let voice_name = file
                .path()
                .file_name()
                .unwrap()
                .to_str()
                .unwrap()
                .to_string();

            if voice_name.eq_ignore_ascii_case(name) {
                return Ok(Some(file.path()));
            }
        }
    }

    Ok(None)
}

#[derive(Debug)]
pub enum VoiceError {
    VoiceNotFound(String),
    WordNotFound(String),
    NotInVoiceChannel,
    Serenity(poise::serenity_prelude::Error),
}

impl Display for VoiceError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            VoiceError::VoiceNotFound(v) => write!(f, "{} voice not found", v),
            VoiceError::WordNotFound(w) => write!(f, "{} is not in dictionary", w),
            VoiceError::Serenity(err) => write!(f, "Serenity error: {}", err),
            VoiceError::NotInVoiceChannel => write!(f, "User not in voice channel"),
        }
    }
}

impl From<poise::serenity_prelude::Error> for VoiceError {
    fn from(value: poise::serenity_prelude::Error) -> Self {
        Self::Serenity(value)
    }
}

impl std::error::Error for VoiceError {}

struct LeaveHandler {
    guild: GuildId,
    manager: Arc<Songbird>,
}

#[async_trait]
impl EventHandler for LeaveHandler {
    async fn act(&self, ctx: &EventContext<'_>) -> Option<Event> {
        if let EventContext::Track(_track) = ctx {
            let has_handler = self.manager.get(self.guild).is_some();

            if has_handler {
                let _ = self.manager.remove(self.guild).await.is_ok();
            }
        }

        None
    }
}

pub async fn speak(
    ctx: Context<'_>,
    guild_id: GuildId,
    user_id: UserId,
    voice: &str,
    phrase: &str,
) -> Result<(), VoiceError> {
    let _ = ctx.data().speak_lock.lock().await;

    let voice_path = match find_voice(&ctx.data().cfg.voice_path, voice).await.unwrap() {
        None => return Err(VoiceError::VoiceNotFound(voice.to_string())),
        Some(voice_path) => voice_path,
    };

    let dict = get_voice_dictionary(&voice_path).await.unwrap();

    let mut sentence = Vec::new();
    for word in phrase.split(' ') {
        let word = word.to_lowercase();
        let mut add_period = false;
        let mut add_comma = false;
        let word = if word.ends_with(',') {
            add_comma = true;
            word.replace(',', "")
        } else if word.ends_with('.') {
            add_period = true;
            word.replace('.', "")
        } else {
            word.to_string()
        };

        if dict.contains_key(&word) {
            sentence.push(word.to_string());
        } else {
            return Err(VoiceError::WordNotFound(word));
        }

        if add_comma {
            sentence.push("_comma".to_string());
        }

        if add_period {
            sentence.push("_period".to_string());
        }
    }

    let channel_id = {
        let guild = ctx.cache().guild(guild_id).unwrap();

        let channel_id = guild
            .voice_states
            .get(&user_id)
            .and_then(|voice_state| voice_state.channel_id);

        channel_id
    };

    let connect_to = match channel_id {
        Some(channel) => channel,
        None => {
            return Err(VoiceError::NotInVoiceChannel);
        }
    };

    let manager = songbird::get(ctx.serenity_context())
        .await
        .expect("Songbird not initialized")
        .clone();

    let handler_lock = manager.join(guild_id, connect_to).await;

    if let Ok(handler_lock) = handler_lock {
        let mut handler = handler_lock.lock().await;
        for (ndx, word) in sentence.iter().enumerate() {
            let word_path = dict.get(word).cloned().unwrap();

            let src = input::File::new(word_path);

            let audio_src = Compressed::new(src.into(), Bitrate::BitsPerSecond(128_000))
                .await
                .expect("Bad params on message load");

            let _ = audio_src.raw.spawn_loader();

            let handle = handler.enqueue_input(audio_src.into()).await;
            handle.set_volume(0.50).unwrap();

            if ndx == sentence.len() - 1 {
                // add task to disconnect bot

                handle
                    .add_event(
                        Event::Track(TrackEvent::End),
                        LeaveHandler {
                            guild: ctx.guild_id().unwrap(),
                            manager: manager.clone(),
                        },
                    )
                    .unwrap();
            }
        }
    }

    Ok(())
}

/// Speak in the language of the gods
#[poise::command(prefix_command, category = "Voice", guild_only)]
pub async fn say(
    ctx: Context<'_>,
    #[description = "Voice to use"] voice: String,
    #[description = "message to say"]
    #[rest]
    phrase: String,
) -> Result<(), Error> {
    let guild_id = ctx.guild_id().unwrap();

    if let Err(err) = speak(ctx, guild_id, ctx.author().id, &voice, &phrase).await {
        match err {
            VoiceError::VoiceNotFound(_)
            | VoiceError::WordNotFound(_)
            | VoiceError::NotInVoiceChannel => {
                ctx.reply(format!("Error: {}", err)).await?;
            }
            _ => return Err(err.into()),
        }
    }

    Ok(())
}

#[poise::command(prefix_command, category = "Voice", guild_only)]
pub async fn list_words(
    ctx: Context<'_>,
    #[description = "Get all words for this voice"] voice: String,
) -> Result<(), Error> {
    match find_voice(&ctx.data().cfg.voice_path, &voice).await? {
        None => {
            ctx.reply(format!("No voice found called '{}'", voice))
                .await?;
        }
        Some(voice_path) => {
            let dict = get_voice_dictionary(&voice_path).await?;
            let mut words: Vec<String> = dict.keys().cloned().collect();

            words.sort();

            let mut list_msg = MessageBuilder::new();
            list_msg.push_line(format!("Here are the words for {}:", voice));

            for word in words {
                list_msg.push("* ");
                list_msg.push_line(word);
            }

            let file_data = list_msg.build();
            let file_data = file_data.as_bytes();
            let file_data = Cow::from(file_data);

            ctx.channel_id()
                .send_message(
                    &ctx,
                    CreateMessage::new()
                        .add_file(CreateAttachment::bytes(file_data, "words.txt".to_string())),
                )
                .await?;
        }
    };

    Ok(())
}

/// List all voices the bot has
#[poise::command(prefix_command, category = "Voice", guild_only)]
pub async fn list_voices(ctx: Context<'_>) -> Result<(), Error> {
    let mut voice_message = MessageBuilder::new();

    for voice in get_voices(&ctx.data().cfg.voice_path).await? {
        voice_message.push("* ");
        voice_message.push_line(voice);
    }

    ctx.reply(voice_message.build()).await?;

    Ok(())
}