diff --git a/src/bot.rs b/src/bot.rs index 7f00b12..6a14e02 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -1,11 +1,12 @@ +use crate::history::MessageHistory; use crate::ExecutorConnection; use async_trait::async_trait; use fancy_regex::{Captures, Regex}; use std::collections::HashMap; -use tokio::sync::{Mutex, RwLock}; +use tokio::sync::Mutex; fn dissect<'a>(prefix: &str, str: &'a str) -> Option<(&'a str, Option<&'a str>)> { - let str = str.strip_prefix(prefix)?; + let str = str.trim().strip_prefix(prefix)?; if let Some(o) = str.find(' ') { Some((&str[..o], Some(&str[o + 1..]))) } else { @@ -15,7 +16,11 @@ fn dissect<'a>(prefix: &str, str: &'a str) -> Option<(&'a str, Option<&'a str>)> #[async_trait] pub trait Trigger { - async fn execute<'a>(&mut self, msg: Context<'a>, captures: Captures<'a>) -> anyhow::Result; + async fn execute<'a>( + &mut self, + msg: Context<'a>, + captures: Captures<'a>, + ) -> anyhow::Result; } #[async_trait] @@ -24,15 +29,15 @@ pub trait Command { } pub struct Context<'a> { - pub last_msg: &'a RwLock>, + pub history: &'a MessageHistory, pub author: &'a str, // in case of triggers, this is always Some(...) pub content: Option<&'a str>, - pub db: &'a ExecutorConnection + pub db: &'a ExecutorConnection, } pub struct Bot anyhow::Result<()>> { - last_msg: RwLock>, + history: MessageHistory, prefix: String, db: ExecutorConnection, commands: HashMap>>, @@ -41,9 +46,9 @@ pub struct Bot anyhow::Result<()>> { } impl anyhow::Result<()>> Bot { - pub fn new(prefix: String, db: ExecutorConnection, sendmsg: SF) -> Self { + pub fn new(prefix: String, db: ExecutorConnection, hdepth: usize, sendmsg: SF) -> Self { Bot { - last_msg: RwLock::new(HashMap::new()), + history: MessageHistory::new(hdepth), commands: HashMap::new(), triggers: Vec::new(), prefix, @@ -69,38 +74,36 @@ impl anyhow::Result<()>> Bot { if let Some((command, remainder)) = dissect(&self.prefix, content) { if let Some(handler) = self.commands.get(command) { let msg = Context { - last_msg: &self.last_msg, author, content: remainder, - db: &self.db + db: &self.db, + history: &self.history, }; - return (self.sendmsg)(origin.into(), handler.lock().await.execute(msg).await?) + return (self.sendmsg)(origin.into(), handler.lock().await.execute(msg).await?); } - return (self.sendmsg)(origin.into(), "Unknown command.".into()) + return (self.sendmsg)(origin.into(), "Unknown command.".into()); } else { for trigger in &self.triggers { let captures = trigger.0.captures(content)?; if let Some(captures) = captures { let msg = Context { - last_msg: &self.last_msg, author, content: Some(content), - db: &self.db + db: &self.db, + history: &self.history, }; - return (self.sendmsg)(origin.into(), trigger.1.lock().await.execute(msg, captures).await?) + return (self.sendmsg)( + origin.into(), + trigger.1.lock().await.execute(msg, captures).await?, + ); } } - self.last_msg.write().await.insert(author.to_string(), content.to_string()); + self.history.add_message(author, content).await; } Ok(()) } - pub async fn handle_message( - &self, - origin: &str, - author: &str, - content: &str, - ) { + pub async fn handle_message(&self, origin: &str, author: &str, content: &str) { if let Err(e) = self.handle_message_inner(origin, author, content).await { let _ = (self.sendmsg)(origin.into(), format!("Error: {}", e)); } diff --git a/src/commands/eval.rs b/src/commands/eval.rs index 714bbc7..50aba13 100644 --- a/src/commands/eval.rs +++ b/src/commands/eval.rs @@ -1,10 +1,10 @@ -use std::collections::HashMap; -use async_trait::async_trait; use crate::bot::{Command, Context}; +use async_trait::async_trait; +use std::collections::HashMap; #[derive(Default)] pub struct Eval { - last_eval: HashMap + last_eval: HashMap, } #[async_trait] diff --git a/src/commands/help.rs b/src/commands/help.rs index e8142af..e7f73c8 100644 --- a/src/commands/help.rs +++ b/src/commands/help.rs @@ -1,4 +1,4 @@ -use crate::bot::{Context, Command}; +use crate::bot::{Command, Context}; use async_trait::async_trait; const HELP: &str = concat!( diff --git a/src/commands/leek.rs b/src/commands/leek.rs index ba4529d..20a9638 100644 --- a/src/commands/leek.rs +++ b/src/commands/leek.rs @@ -113,11 +113,11 @@ enum LeekCommand { async fn execute_leek(cmd: LeekCommand, msg: &Context<'_>) -> anyhow::Result { let nick = msg.content.unwrap_or(msg.author); - match msg.last_msg.read().await.get(nick) { + match msg.history.last_msg(nick).await { Some(msg) => Ok(match cmd { - LeekCommand::Owo => owoify(msg)?, - LeekCommand::Leet => leetify(msg), - LeekCommand::Mock => mock(msg), + LeekCommand::Owo => owoify(&msg)?, + LeekCommand::Leet => leetify(&msg), + LeekCommand::Mock => mock(&msg), } .to_string()), None => Ok("No previous messages found.".into()), diff --git a/src/commands/mod.rs b/src/commands/mod.rs index 363279e..dbef554 100644 --- a/src/commands/mod.rs +++ b/src/commands/mod.rs @@ -1,10 +1,11 @@ #[cfg(feature = "debug")] pub mod debug; +pub mod eval; pub mod help; pub mod leek; -pub mod waifu; +pub mod quotes; pub mod sed; -pub mod eval; pub mod spotify; -pub mod title; \ No newline at end of file +pub mod title; +pub mod waifu; diff --git a/src/commands/quotes.rs b/src/commands/quotes.rs new file mode 100644 index 0000000..ad637f9 --- /dev/null +++ b/src/commands/quotes.rs @@ -0,0 +1,58 @@ +use crate::bot::{Command, Context}; +use crate::database::Quote; +use async_trait::async_trait; + +pub struct Grab; +pub struct Quot; + +#[async_trait] +impl Command for Grab { + async fn execute(&mut self, msg: Context<'_>) -> anyhow::Result { + let content = if let Some(c) = msg.content { + c + } else { + return Ok("Invalid usage.".into()); + }; + let mut split = content.splitn(2, ' '); + let split = (split.next().unwrap(), split.next()); + let (author, count) = if let Some(author) = split.1 { + (author, split.0.parse::()?) + } else { + (split.0, 1) + }; + if count == 0 { + return Ok("So are you going to grab anything?".into()); + } + if author == msg.author { + return Ok("You can't grab yourself.".into()); + } + let message = msg + .history + .last_msgs(author, count) + .await + .map(|v| v.join(" | ")); + if let Some(message) = message { + msg.db + .add_quote(Quote { + author: author.into(), + quote: message, + }) + .await?; + Ok("Quote added ({} messages).".into()) + } else { + Ok("No previous messages to grab.".into()) + } + } +} + +#[async_trait] +impl Command for Quot { + async fn execute(&mut self, msg: Context<'_>) -> anyhow::Result { + let author = msg.content.map(ToString::to_string); + if let Some(q) = msg.db.get_quote(author).await? { + Ok(format!("\"{}\" ~{}", q.quote, q.author)) + } else { + Ok("No quotes found from this user.".into()) + } + } +} diff --git a/src/commands/sed.rs b/src/commands/sed.rs index 2f7be13..ab124f8 100644 --- a/src/commands/sed.rs +++ b/src/commands/sed.rs @@ -1,12 +1,16 @@ +use crate::bot::{Context, Trigger}; use async_trait::async_trait; use fancy_regex::Captures; -use crate::bot::{Context, Trigger}; pub struct Sed; #[async_trait] impl Trigger for Sed { - async fn execute<'a>(&mut self, msg: Context<'a>, captures: Captures<'a>) -> anyhow::Result { + async fn execute<'a>( + &mut self, + msg: Context<'a>, + captures: Captures<'a>, + ) -> anyhow::Result { let foreign_author; let author = if let Some(author) = captures.name("u").map(|m| m.as_str()) { foreign_author = true; @@ -15,8 +19,7 @@ impl Trigger for Sed { foreign_author = false; msg.author }; - let lastmsg = msg.last_msg.read().await; - let message = if let Some(msg) = lastmsg.get(author) { + let message = if let Some(msg) = msg.history.last_msg(author).await { msg } else { return Ok("No previous messages found.".into()); @@ -25,11 +28,15 @@ impl Trigger for Sed { // TODO: karx plz add flags //let flags = matches.name("f").map(|m| m.as_str()); let result = message.replace(find.as_str(), replace.as_str()); - drop(lastmsg); if foreign_author { - Ok(format!("(edited by {}) <{}> {}", msg.author, author, result)) + Ok(format!( + "(edited by {}) <{}> {}", + msg.author, author, result + )) } else { - msg.last_msg.write().await.insert(author.into(), result.to_string()); + msg.history + .edit_message(author, 0, result.to_string()) + .await; Ok(format!("<{}> {}", author, result)) } } else { diff --git a/src/history.rs b/src/history.rs index 45c4d04..beaa562 100644 --- a/src/history.rs +++ b/src/history.rs @@ -1,24 +1,19 @@ use std::collections::{HashMap, VecDeque}; -use anyhow::anyhow; use tokio::sync::RwLock; pub struct MessageHistory { map: RwLock>>, - maxlen: usize + maxlen: usize, } impl MessageHistory { pub fn new(maxlen: usize) -> MessageHistory { MessageHistory { map: RwLock::new(HashMap::new()), - maxlen + maxlen, } } - pub fn max_len(&self) -> usize { - self.maxlen - } - pub async fn last_msg(&self, user: &str) -> Option { let map = self.map.read().await; map.get(user).and_then(|d| d.get(0)).map(|s| s.to_string()) @@ -27,18 +22,39 @@ impl MessageHistory { pub async fn last_msgs(&self, user: &str, count: usize) -> Option> { let map = self.map.read().await; if let Some(deque) = map.get(user) { - let count = if deque.len() < count { deque.len() } else {count}; - Some(deque.range(..count).rev().map(ToString::to_string).collect()) + let count = if deque.len() < count { + deque.len() + } else { + count + }; + Some( + deque + .range(..count) + .rev() + .map(ToString::to_string) + .collect(), + ) } else { None } } + pub async fn edit_message(&self, user: &str, pos: usize, edited: String) -> bool { + let mut map = self.map.write().await; + if let Some(deque) = map.get_mut(user) { + if let Some(old) = deque.get_mut(pos) { + *old = edited; + return true; + } + } + false + } + pub async fn add_message(&self, user: &str, message: &str) { let mut map = self.map.write().await; if let Some(deque) = map.get_mut(user) { if deque.len() == self.maxlen { - deque.remove(deque.len()-1); + deque.remove(deque.len() - 1); } deque.push_front(message.to_string()); } else { @@ -47,4 +63,4 @@ impl MessageHistory { map.insert(user.to_string(), deque); } } -} \ No newline at end of file +} diff --git a/src/main.rs b/src/main.rs index 0c640fd..8b5b558 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,11 +1,18 @@ +use fancy_regex::Regex; use std::env; use std::fs::File; use std::io::Read; use std::sync::Arc; use std::thread; -use fancy_regex::Regex; use crate::bot::Bot; +use crate::commands::eval::Eval; +use crate::commands::help::Help; +use crate::commands::leek::Owo; +use crate::commands::quotes::{Grab, Quot}; +use crate::commands::sed::Sed; +use crate::commands::spotify::Spotify; +use crate::commands::title::Title; use crate::commands::waifu::Waifu; use futures_util::stream::StreamExt; use irc::client::prelude::Config; @@ -16,12 +23,6 @@ use tokio::select; use tokio::sync::broadcast; use tokio::sync::mpsc::unbounded_channel; use tracing_subscriber::EnvFilter; -use crate::commands::eval::Eval; -use crate::commands::help::Help; -use crate::commands::leek::Owo; -use crate::commands::sed::Sed; -use crate::commands::spotify::Spotify; -use crate::commands::title::Title; use crate::config::UberConfig; use crate::database::{DbExecutor, ExecutorConnection}; @@ -30,6 +31,7 @@ mod bot; mod commands; mod config; mod database; +mod history; #[cfg(unix)] async fn terminate_signal() { @@ -89,7 +91,7 @@ async fn main() -> anyhow::Result<()> { let (ctx, _) = broadcast::channel(1); let (etx, mut erx) = unbounded_channel(); - let mut bot = Bot::new(cfg.irc.prefix, db_conn, { + let mut bot = Bot::new(cfg.irc.prefix, db_conn, 3, { let client = client.clone(); move |target, msg| Ok(client.send_privmsg(target, msg)?) }); @@ -98,7 +100,12 @@ async fn main() -> anyhow::Result<()> { bot.add_command("waifu".into(), Waifu::default()); bot.add_command("owo".into(), Owo); bot.add_command("ev".into(), Eval::default()); - bot.add_trigger(Regex::new(r"^(?:(?\S+):\s+)?s/(?[^/]*)/(?[^/]*)(?:/(?[a-z]*))?\s*")?, Sed); + bot.add_command("grab".into(), Grab); + bot.add_command("quot".into(), Quot); + bot.add_trigger( + Regex::new(r"^(?:(?\S+):\s+)?s/(?[^/]*)/(?[^/]*)(?:/(?[a-z]*))?\s*")?, + Sed, + ); if let Some(spotcfg) = cfg.spotify { let creds = Credentials::new(&spotcfg.client_id, &spotcfg.client_secret); let spotify = Spotify::new(creds).await?; @@ -150,7 +157,7 @@ async fn main() -> anyhow::Result<()> { async fn message_loop anyhow::Result<()>>( mut stream: ClientStream, - bot: Bot + bot: Bot, ) -> anyhow::Result<()> { while let Some(message) = stream.next().await.transpose()? { if let Command::PRIVMSG(ref origin, content) = message.command {