FrenBot/src/discord/story.rs

use crate::config::{BotConfig, Channel};
use crate::{command, group, GlobalData};
use rand::prelude::{IndexedRandom, SliceRandom};
use rand::thread_rng;
use regex::Regex;
use serenity::client::Context;
use serenity::framework::standard::{Args, CommandResult};
use serenity::model::channel::{Attachment, Message};
use serenity::utils::MessageBuilder;
use std::collections::HashMap;
use std::path::PathBuf;

#[group]
#[commands(story, word, list_stories, upload_story)]
pub struct Story;

async fn get_all_stories(cfg: &BotConfig) -> Vec<PathBuf> {
    let mut dir = tokio::fs::read_dir(&cfg.story_path).await.unwrap();

    let mut stories = Vec::new();

    while let Some(file) = &dir.next_entry().await.unwrap() {
        stories.push(file.path());
    }

    stories
}

fn get_all_blanks(story: &str) -> Vec<String> {
    let re = Regex::new(r"\{%(?P<global>.*?)%}").unwrap();
    let mut globals: Vec<String> = Vec::new();

    for cap in re.captures_iter(story) {
        if let Some(global) = cap.name("global") {
            if !globals.contains(&global.as_str().to_string()) {
                globals.push(global.as_str().to_string())
            }
        }
    }

    globals
}

#[command]
#[only_in(guilds)]
#[description("List Stories")]
async fn list_stories(ctx: &Context, msg: &Message, _args: Args) -> CommandResult {
    let data = ctx.data.read().await;
    let global_data = data.get::<GlobalData>().unwrap();

    let stories = get_all_stories(&global_data.cfg).await;

    let mut resp = MessageBuilder::default();

    resp.push_line("**Stories I know:**");

    for story in stories {
        resp.push_line(
            story
                .file_stem()
                .unwrap()
                .to_str()
                .unwrap()
                .replace('_', " "),
        );
    }

    msg.reply(&ctx.http, resp.build()).await?;

    Ok(())
}

#[command]
#[only_in(guilds)]
#[description("Let me tell you a tail")]
async fn story(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
    let story_channel = {
        let data = ctx.data.read().await;
        data.get::<Channel<String>>().unwrap().clone()
    };

    let mut story_recv = match story_channel.recv.try_lock() {
        Ok(story_recv) => story_recv,
        Err(_) => {
            msg.reply(&ctx.http, "Let me finish this story bub").await?;
            return Ok(());
        }
    };

    let stories = {
        let data = ctx.data.read().await;

        let global_data = data.get::<GlobalData>().unwrap();
        get_all_stories(&global_data.cfg).await
    };

    let story_path = if args.is_empty() {
        let mut rng = thread_rng();

        stories.choose(&mut rng)
    } else {
        let story_name = args.rest().to_lowercase().replace(' ', "_");

        stories.iter().find(|story_p| {
            let story_p_name = story_p.file_stem().unwrap().to_str().unwrap();
            story_p_name == story_name
        })
    };

    let story_path = match story_path {
        None => {
            msg.reply(&ctx.http, "No story found :(").await?;
            return Ok(());
        }
        Some(story) => story,
    };

    let mut story_contents = tokio::fs::read_to_string(story_path).await?;

    let mut story_globals: HashMap<String, String> = HashMap::new();

    let msg_channel = msg.channel(&ctx.http).await?.id();

    for global in get_all_blanks(&story_contents) {
        msg_channel
            .say(&ctx.http, format!("Give me {}", global))
            .await?;

        story_globals.insert(global, story_recv.recv().await.unwrap());
    }

    for (prompt, response) in story_globals {
        story_contents = story_contents.replace(&format!("{{%{}%}}", prompt), &response)
    }

    let mut msg_builder = MessageBuilder::default();
    for part in story_contents.split('\n') {
        if (msg_builder.0.len() + part.len()) > serenity::constants::MESSAGE_CODE_LIMIT {
            msg_channel.say(&ctx.http, msg_builder.build()).await?;
            msg_builder = MessageBuilder::default();
        }

        msg_builder.push_line(part);
    }

    if !msg_builder.0.is_empty() {
        msg_channel.say(&ctx.http, msg_builder.build()).await?;
    }

    Ok(())
}

#[command]
#[only_in(guilds)]
#[description("Give me a word")]
async fn word(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
    let data = ctx.data.read().await;
    let story_channel = data.get::<Channel<String>>().unwrap();

    if story_channel.recv.try_lock().is_ok() {
        msg.reply(&ctx.http, "No story in progress!").await?;
        return Ok(());
    }

    let story_send = story_channel.send.lock().await;

    let resp = MessageBuilder::default().push_safe(args.rest()).build();

    story_send.send(resp).await?;

    Ok(())
}

#[command]
#[description("Attach stories as .txt files")]
async fn upload_story(ctx: &Context, msg: &Message, _args: Args) -> CommandResult {
    let data = ctx.data.read().await;
    let global = data.get::<GlobalData>().unwrap();

    let stories: Vec<&Attachment> = msg
        .attachments
        .iter()
        .filter(|f| f.filename.ends_with(".txt"))
        .collect();

    for story in stories {
        let story_txt = String::from_utf8(story.download().await?)?;

        let globals = get_all_blanks(&story_txt);

        let mut resp = MessageBuilder::new();

        resp.push_line(format!("{} contains the following blanks:", story.filename));

        for blank in globals {
            resp.push_line(format!("* {}", blank));
        }

        msg.reply(&ctx.http, resp.build()).await?;

        let story_path = global.cfg.story_path.join(&story.filename);

        tokio::fs::write(&story_path, story_txt).await?;
    }

    Ok(())
}