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, GuildId},
        user::User,
    },
    prelude::*,
};
use std::fs::read_to_string;
use std::sync::Once;
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 Handler {
    start_loop: Once,
}

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: 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::sleep(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);
    }

    async fn cache_ready(&self, ctx: Context, _guilds: Vec<GuildId>) {
        self.start_loop.call_once(move || {
            tokio::spawn(async move { interval_loop(ctx).await });
        });
    }
}

#[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::builder(&config.token)
        .event_handler(Handler {
            start_loop: Once::new(),
        })
        .type_map_insert::<TextGenerator>(generator)
        .type_map_insert::<Interval>(config.interval)
        .type_map_insert::<ChanId>(config.channel_id)
        .await
        .expect("err creating client");

    client.start().await.unwrap();
}