generator-bot/src/main.rs

extern crate markov;
extern crate purrchance;
extern crate serde;
extern crate serenity;
extern crate tokio;
extern crate toml;
// extern crate tracery;

use purrchance::Purrchance;
use serde::Deserialize;
use serenity::{
    async_trait,
    model::{channel::Message, gateway::Ready, id::ChannelId, user::User},
    prelude::*,
};
use std::fs::read_to_string;
use std::sync::Arc;
use std::time::Duration;

enum TextGenerator {
    Perchance { grammar: purrchance::Grammar, symbol: purrchance::Symbol },
//    Tracery(tracery::Grammar),
    Markov(markov::Chain<String>),
}

impl TextGenerator {
    fn generate(&self) -> String {
        match self {
            TextGenerator::Perchance { grammar, symbol } => symbol.eval(&grammar).unwrap(),
//            TextGenerator::Tracery(grammar) => grammar.flatten(),
            TextGenerator::Markov(chain) => chain.generate_str(),
        }
    }
}

#[derive(Deserialize)]
#[serde(tag = "type", rename_all = "lowercase")]
enum GeneratorSpec {
    Perchance { path: String, symbol: String },
//    Tracery { path: String, rule: String },
    Markov { path: String, order: usize },
}

impl GeneratorSpec {
    fn load(&self) -> TextGenerator {
        match self {
            GeneratorSpec::Perchance { path: p, symbol: s } => {
                let grammar = purrchance::parser::load_grammar(&read_to_string(p).unwrap()).unwrap();
                let symbol = purrchance::Symbol::NonTerminal(String::from(s));
                TextGenerator::Perchance { grammar, symbol }
            }
/*
            GeneratorSpec::Tracery { path, rule } => {
                let mut grammar = tracery::from_json(read_to_string(path).unwrap()).unwrap();
                grammar.default_rule(rule);
                TextGenerator::Tracery(grammar)
            }
*/
            GeneratorSpec::Markov { path, order } => {
                let mut chain = markov::Chain::of_order(*order);
                for line in read_to_string(path).unwrap().lines() {
                    chain.feed_str(line);
                }
                TextGenerator::Markov(chain)
            }
        }
    }
}

#[derive(Deserialize)]
struct Config {
    generator: GeneratorSpec,
    token: String,
    channel_id: ChannelId,
    interval: u64,
}


impl TypeMapKey for TextGenerator {
    type Value = TextGenerator;
}

struct Interval;

impl TypeMapKey for Interval {
    type Value = u64;
}

struct ChanId;

impl TypeMapKey for ChanId {
    type Value = ChannelId;
}

struct LoopStatus;

impl TypeMapKey for LoopStatus {
    type Value = bool;
}

struct Handler;

async fn mentions_me(ctx: &Context, msg: &Message) -> bool {
    let me = User::from(ctx.http.get_current_user().await.unwrap());
    if msg.mentions_user(&me) {
        return true
    }
    let guild = msg.guild(ctx).await.unwrap();
    let member = guild.member(ctx, me).await.unwrap();
    let my_roles = member.roles;
    msg.mention_roles.iter().any(|r| my_roles.contains(r))
}

async fn interval_loop(ctx: Arc<Context>) {
    let channel_id = ctx.data.read().await.get::<ChanId>().unwrap().clone();
    loop {
        let sentence = ctx.data.read().await.get::<TextGenerator>().unwrap().generate();
        channel_id.say(&ctx.http, sentence).await.unwrap();
        tokio::time::delay_for(Duration::from_secs(*ctx.data.read().await.get::<Interval>().unwrap())).await;
    }
}

#[async_trait]
impl EventHandler for Handler {
    async fn message(&self, ctx: Context, msg: Message) {
        if mentions_me(&ctx, &msg).await {
            let sentence = ctx.data.read().await.get::<TextGenerator>().unwrap().generate();
            msg.channel_id.say(&ctx.http, sentence).await.unwrap();
        }
    }

    async fn ready(&self, ctx: Context, ready: Ready) {
        eprintln!("{} is connected!", ready.user.name);
        let ctx = Arc::new(ctx);
        let ctx_clone = Arc::clone(&ctx);
        if !ctx.data.read().await.get::<LoopStatus>().unwrap() {
            let interval_loop = tokio::spawn(async move {interval_loop(ctx_clone).await});
            {
                let mut data = ctx.data.write().await;
                data.insert::<LoopStatus>(true);
            }
            let _ = interval_loop.await;
            {
                let mut data = ctx.data.write().await;
                data.insert::<LoopStatus>(false);
            }
        }
    }
}

#[tokio::main]
async fn main() {
    let config: Config = toml::from_str(&read_to_string("config.toml").unwrap()).unwrap();

    let generator = config.generator.load();

    let mut client = Client::new(&config.token)
        .event_handler(Handler)
        .await
        .expect("err creating client");

    {
        let mut data = client.data.write().await;
        data.insert::<TextGenerator>(generator);
        data.insert::<Interval>(config.interval);
        data.insert::<ChanId>(config.channel_id);
        data.insert::<LoopStatus>(false);
    }
    client.start().await.unwrap();
}