Skip to content
Snippets Groups Projects
handler.rs 10.4 KiB
Newer Older
Stephen D's avatar
Stephen D committed
use crate::framebuffer::FrameBuffer;
use crate::program::Program;
use diesel::PgConnection;
use phf::phf_map;
use serenity::async_trait;
use serenity::http::AttachmentType;
use serenity::model::channel::{Message, ReactionType};
use serenity::model::id::UserId;
use serenity::model::prelude::Ready;
use serenity::prelude::*;
use std::borrow::{Borrow, Cow};
Stephen D's avatar
Stephen D committed
use std::collections::HashMap;
use std::str::FromStr;
use std::sync::{Arc, Mutex};
Stephen D's avatar
Stephen D committed
use tokio::task;
Stephen D's avatar
Stephen D committed
use xbasic::basic_io::BasicIO;
use xbasic::expr::ExprValue;
use xbasic::xbasic::XBasicBuilder;

static EMOJI_MAP: phf::Map<&'static str, &'static str> = phf_map! {
"cat" => "🐈",
"chicken" => "🐔",
"spaghetti" => "🍝",
"dog" => "🐕",
"bot" => "🤖",
"mango" => "🥭",
"banana" => "🍌",
"bee" => "🐝"
};

struct DiscordIO {
	s: String,
	frame: Option<FrameBuffer>,
}

impl DiscordIO {
	fn new() -> Self {
		Self {
			s: String::new(),
			frame: None,
		}
	}
}

impl BasicIO for DiscordIO {
	fn read_line(&mut self) -> String {
		unimplemented!()
	}

	fn write_line(&mut self, line: String) {
		self.s += &*(line + "\r\n");
	}
}

Stephen D's avatar
Stephen D committed
macro_rules! get_user_programs {
	($self: expr, $user_id: expr) => {
		$self.programs.lock().unwrap().get_mut($user_id).unwrap()
	};
}

Stephen D's avatar
Stephen D committed
pub(crate) struct Handler {
	programs: Arc<Mutex<HashMap<UserId, Program>>>,
	conn: Arc<Mutex<PgConnection>>,
}

impl Handler {
	pub fn new(conn: PgConnection) -> Self {
		Self {
			programs: Arc::new(Mutex::new(HashMap::new())),
			conn: Arc::new(Mutex::new(conn)),
		}
	}

Stephen D's avatar
Stephen D committed
	async fn interpret_line(&self, msg: &Message, ctx: &Context, line: &str) {
		// TODO we lock the mutex to check, but unlock before locking again later
		// allows another thread to screw it up
		// we should lock once for this entire function
		if self.programs.lock().unwrap().contains_key(&msg.author.id) {
			match line {
				"!STOP" => {
					self.interpreter_stop(msg);
				}
				"RUN" => {
					self.interpreter_run(msg, ctx).await;
				}
				"LIST" => {
					self.interpreter_list(msg, ctx).await;
				}
				"DIR" => {
					if self.list_saved_programs(msg, ctx).await.is_none() {
						msg.channel_id
Stephen D's avatar
Stephen D committed
							.say(&ctx, "Could not get list of programs from database.")
							.await
							.unwrap();
					}
				}
				"PUBDIR" => {
					if self.list_published_programs(msg, ctx).await.is_none() {
						msg.channel_id
							.say(
								&ctx,
								"Could not get list of published programs from database.",
							)
							.await
							.unwrap();
Stephen D's avatar
Stephen D committed
				_ => {
					if let Some(name) = line.strip_prefix("SAVE ") {
						self.interpreter_save(name, msg, ctx).await;
					}

					if let Some(name) = line.strip_prefix("LOAD ") {
						self.interpreter_load(name, msg, ctx).await;
					}

Stephen D's avatar
Stephen D committed
					if let Some(name) = line.strip_prefix("PUB ") {
						if self.publish_program(name, msg, ctx).await.is_none() {
							msg.channel_id
								.say(&ctx, format!("Could not publish {}.", name))
								.await
								.unwrap();
						}
					}

					if let Some(name) = line.strip_prefix("UNPUB ") {
						if self.unpublish_program(name, msg, ctx).await.is_none() {
							msg.channel_id
								.say(&ctx, &format!("Could not unpublish {}.", name))
								.await
								.unwrap();
						}
					}

					if let Some(id_str) = line.strip_prefix("PUBLOAD ") {
						match id_str.parse::<i32>() {
							Ok(id) => {
								if self.load_published_program(msg, ctx, id).await.is_none() {
									msg.channel_id
Stephen D's avatar
Stephen D committed
										.say(&ctx, format!("Could not load {}.", id))
Stephen D's avatar
Stephen D committed
										.await
										.unwrap();
								}
							}
							Err(_) => {
								msg.channel_id
									.say(&ctx, "PUBLOAD requires a numerical ID.")
									.await
									.unwrap();
							}
						}
					}

Stephen D's avatar
Stephen D committed
					let mut split = line.splitn(2, ' ');
					let first = split.next().unwrap();
					if let Ok(num) = first.parse::<u32>() {
						match split.next() {
							Some(x) => {
								if x.is_empty() {
									let _ =
										get_user_programs!(&self, &msg.author.id).code.remove(&num);
									return;
								}

								get_user_programs!(&self, &msg.author.id)
									.code
									.insert(num, x.to_owned());
							}
							None => {
								let _ = get_user_programs!(&self, &msg.author.id).code.remove(&num);
							}
						}
Stephen D's avatar
Stephen D committed
		if line == "!START" {
			self.interpreter_start(msg);
		}
	}
Stephen D's avatar
Stephen D committed

Stephen D's avatar
Stephen D committed
	fn interpreter_stop(&self, msg: &Message) {
		self.programs
			.lock()
			.unwrap()
			.remove(&msg.author.id)
			.unwrap();
	}
Stephen D's avatar
Stephen D committed

Stephen D's avatar
Stephen D committed
	fn interpreter_start(&self, msg: &Message) {
		self.programs
			.lock()
			.unwrap()
			.insert(msg.author.id, Program::new());
	}
Stephen D's avatar
Stephen D committed

	async fn list_saved_programs(&self, msg: &Message, ctx: &Context) -> Option<()> {
		let program_names =
			Program::list_programs_by_user(self.conn.lock().ok()?.borrow(), msg.author.id)?;
		msg.channel_id
			.say(
				&ctx,
				format!(
					"You have {} programs saved:\r\n```\r\n{}\r\n```",
					program_names.len(),
					program_names.join("\r\n")
				),
			)
			.await
			.ok()?;

		Some(())
	}

Stephen D's avatar
Stephen D committed
	async fn list_published_programs(&self, msg: &Message, ctx: &Context) -> Option<()> {
		let program_names: Vec<String> =
			Program::list_published_programs(self.conn.lock().ok()?.borrow())?
				.iter()
				.map(|row| format!("{}\t{}", row.0, row.1))
				.collect();

		msg.channel_id
			.say(
				&ctx,
				format!(
				"Found {} public programs:\r\n```\r\nid\tname\r\n{}\r\n```\r\n Load one with `PUBLOAD <id>`",
					program_names.len(),
				program_names.join("\r\n")),
			)
			.await
			.unwrap();

		Some(())
	}

	async fn publish_program(&self, name: &str, msg: &Message, ctx: &Context) -> Option<()> {
		Program::set_program_published(self.conn.lock().ok()?.borrow(), name, msg.author.id, true)?;

		msg.channel_id
			.say(&ctx, format!("Published {}.", name))
			.await
			.unwrap();

		Some(())
	}

	async fn unpublish_program(&self, name: &str, msg: &Message, ctx: &Context) -> Option<()> {
		Program::set_program_published(
			self.conn.lock().ok()?.borrow(),
			name,
			msg.author.id,
			false,
		)?;

		msg.channel_id
Stephen D's avatar
Stephen D committed
			.say(&ctx, format!("Unpublished {}.", name))
Stephen D's avatar
Stephen D committed
			.await
			.unwrap();

		Some(())
	}

	async fn load_published_program(&self, msg: &Message, ctx: &Context, id: i32) -> Option<()> {
		let name = get_user_programs!(&self, &msg.author.id)
			.load_published_program(&self.conn.lock().ok()?.borrow(), id)?;

		msg.channel_id
			.say(&ctx, format!("Loaded {} (\"{}\") into memory.", id, name))
			.await
			.unwrap();

		Some(())
	}

Stephen D's avatar
Stephen D committed
	async fn interpreter_list(&self, msg: &Message, ctx: &Context) {
		msg.channel_id
			.say(
				&ctx,
				format!(
					"```\n{}\n```",
					self.programs.lock().unwrap()[&msg.author.id].stringy_line_nums()
				),
			)
			.await
			.unwrap();
	}
Stephen D's avatar
Stephen D committed

Stephen D's avatar
Stephen D committed
	async fn interpreter_load(&self, name: &str, msg: &Message, ctx: &Context) {
		let result = get_user_programs!(&self, &msg.author.id).load_program(
			&self.conn.lock().unwrap(),
			msg.author.id,
			name,
		);
		match result {
			Some(_) => {
				msg.channel_id
					.say(&ctx, format!("Loaded {} into memory.", name))
					.await
					.unwrap();
			}
			None => {
				msg.channel_id
					.say(&ctx, "Could not load program into memory.")
					.await
					.unwrap();
			}
		}
	}
Stephen D's avatar
Stephen D committed

Stephen D's avatar
Stephen D committed
	async fn interpreter_save(&self, name: &str, msg: &Message, ctx: &Context) {
		let result = get_user_programs!(&self, &msg.author.id).save_program(
			&self.conn.lock().unwrap(),
			msg.author.id,
			name,
		);
		match result {
			Some(_) => {
				msg.channel_id
					.say(&ctx, format!("Saved as {}", name))
					.await
					.unwrap();
			}
			None => {
				msg.channel_id
					.say(&ctx, "Could not save program.")
					.await
					.unwrap();
			}
		}
	}
Stephen D's avatar
Stephen D committed

Stephen D's avatar
Stephen D committed
	async fn interpreter_run(&self, msg: &Message, ctx: &Context) {
		let thinking_react = match ReactionType::from_str("🤔") {
			Ok(react) => msg.react(&ctx, react).await.ok(),
			Err(_) => None,
		};
Stephen D's avatar
Stephen D committed

Stephen D's avatar
Stephen D committed
		let code = self.programs.lock().unwrap()[&msg.author.id].stringify();
		let io = DiscordIO::new();
Stephen D's avatar
Stephen D committed

Stephen D's avatar
Stephen D committed
		let (output, fb, errors) = task::spawn_blocking(move || {
Stephen D's avatar
Stephen D committed
			let mut xbb = XBasicBuilder::new(io);
			xbb.compute_limit(1000000000);
Stephen D's avatar
Stephen D committed

Stephen D's avatar
Stephen D committed
			xbb.define_function("setframe".to_owned(), 2, |args, io| {
				let w = args[0].clone().into_decimal() as u32;
				let h = args[1].clone().into_decimal() as u32;
Stephen D's avatar
Stephen D committed

Stephen D's avatar
Stephen D committed
				io.frame = Some(FrameBuffer::new(w, h));
Stephen D's avatar
Stephen D committed
				ExprValue::Decimal(0.0)
			})
			.unwrap();
Stephen D's avatar
Stephen D committed

Stephen D's avatar
Stephen D committed
			xbb.define_function("setpixel".to_owned(), 5, |args, io| {
				let x = args[0].clone().into_decimal() as u32;
				let y = args[1].clone().into_decimal() as u32;
				let red = args[2].clone().into_decimal() as u8;
				let green = args[3].clone().into_decimal() as u8;
				let blue = args[4].clone().into_decimal() as u8;
Stephen D's avatar
Stephen D committed

Stephen D's avatar
Stephen D committed
				match &mut io.frame {
					Some(fb) => {
						fb.set_pixel(x, y, red, green, blue, 255);
Stephen D's avatar
Stephen D committed
					}
Stephen D's avatar
Stephen D committed
					None => {}
Stephen D's avatar
Stephen D committed
				}
Stephen D's avatar
Stephen D committed

				ExprValue::Decimal(0.0)
			})
			.unwrap();

			let mut xb = xbb.build();

			let _ = xb.run(&format!("{}\n", code));

			let errors = if xb.error_handler.had_errors || xb.error_handler.had_runtime_error {
				Some(xb.error_handler.errors.join("\n"))
			} else {
				None
			};

			(xb.get_io().s.clone(), xb.get_io().frame.clone(), errors)
Stephen D's avatar
Stephen D committed
		})
		.await
		.unwrap();
Stephen D's avatar
Stephen D committed

		if !output.is_empty() {
			msg.channel_id.say(&ctx, output).await.unwrap();
		}

		if let Some(fb) = &fb {
			let buf = fb.as_png_vec();

			msg.channel_id
				.send_message(&ctx, |e| {
					e.add_file(AttachmentType::Bytes {
						data: Cow::Borrowed(&buf),
						filename: "output.png".to_string(),
					});

					e
				})
				.await
				.unwrap();
		}

		if let Some(e) = errors {
			msg.channel_id.say(&ctx, e).await.unwrap();
			if let Ok(react) = ReactionType::from_str("❌") {
				let _ = msg.react(&ctx, react).await;
Stephen D's avatar
Stephen D committed
			}
Stephen D's avatar
Stephen D committed
		} else if let Ok(react) = ReactionType::from_str("✅") {
			let _ = msg.react(&ctx, react).await;
		}
Stephen D's avatar
Stephen D committed

Stephen D's avatar
Stephen D committed
		if let Some(t) = thinking_react {
			let _ = t.delete(&ctx).await;
		}
	}
}

#[async_trait]
impl EventHandler for Handler {
	async fn message(&self, ctx: Context, msg: Message) {
		for (key, value) in EMOJI_MAP.entries() {
			let msg_lower = format!(" {} ", msg.content.to_lowercase());
			if msg_lower.contains(&format!(" {} ", key))
				|| msg_lower.contains(&format!(" {}s ", key))
			{
				let reaction_type = match ReactionType::from_str(value) {
					Ok(x) => x,
					Err(x) => {
						println!("Could not react: {}", x);
						return;
					}
				};
				if let Err(e) = msg.react(&ctx, reaction_type).await {
					println!("Error reacting: {}", e);
				}
Stephen D's avatar
Stephen D committed
			}
		}
Stephen D's avatar
Stephen D committed

		for line in msg.content.split('\n') {
			self.interpret_line(&msg, &ctx, line).await;
		}
Stephen D's avatar
Stephen D committed
	}

	async fn ready(&self, _: Context, ready: Ready) {
		println!("{} is connected", ready.user.name);
	}
}