use std::ops::{Deref, DerefMut}; use std::time::Duration; use scylla::client::session::Session; use crate::database::models::{Channel, Message, User}; use crate::extras::Extras; use crate::framework::commands::arguments::{ArgumentParsingError, FromArguments, NeedyString}; use crate::framework::commands::context::CommandContext; use crate::state::State; use crate::utils::split_at_nth_from_end; /// Fetch a user by their ID from the database. /// /// Arguments: /// * `id` - The user's ID. /// * `db` - The database session. /// /// Returns: /// [`Result`] - The user from the DB, or an [`ArgumentParsingError`] ready /// to be returned by the parser. async fn get_user_by_id(id: u64, db: &Session) -> Result { let user = User::get_by_id(id, db).await?; match user { Some(user) => Ok(user), None => Err(ArgumentParsingError::InvalidArgument), } } /// Fetches the user from the message's reply and falls back to the command executor. /// /// This function takes into account messages on call and traces them back to the original author. /// /// Arguments: /// * `ctx` - The command context. /// /// Returns: /// [`Result`] - The user from the message's reply, or an /// [`ArgumentParsingError`] ready to be returned by the parser. async fn user_from_reply_or_command_executor( ctx: &CommandContext<'_, State, Extras>, ) -> Result { if let Some(referenced) = &ctx.event.referenced_message && !referenced.author.system.unwrap_or(false) { if referenced.author.bot && let Some(message) = Message::get_by_propagation_message_id(referenced.id.get(), &ctx.rt.state.db) .await? { return get_user_by_id(message.author_id, &ctx.rt.state.db).await; } // If the replied message's author is not a bot or system user, use them. return get_user_by_id(referenced.author.id.get(), &ctx.rt.state.db).await; } // fallback to command executor get_user_by_id(ctx.event.author.id.get(), &ctx.rt.state.db).await } /// Parses a user mention or ID into a UserArgument. /// /// It parses from multiple sources, attempting in the following way: /// /// 1. If the argument is a Discord mention, parsing it. /// 2. If the argument is a Discord ID, finding it in the DB. /// 3. If the argument is a Discord mention, finding it in the DB. /// 4. If there's no argument and the message replies to another message: /// a. If the replied message's author is a bot, attempts to query the user who sent it from the /// other side of the call. If it fails, the command executor is used. /// b. If the replied message's author is not a bot, they're used. /// 5. If there's no argument and the message is not replying to any message, the command executor /// is used. pub struct UserArgument(pub User); impl FromArguments for UserArgument { async fn from_args<'a>( args: &'a str, ctx: &'a CommandContext<'_, State, Extras>, ) -> Result<(Self, &'a str), ArgumentParsingError> { if !args.is_empty() { // There's an argument, so that argument is used to find the user. let user: User; let (identifier, args) = NeedyString::from_args(args, ctx).await?; let identifier = identifier.0.trim(); if identifier.starts_with("<@") && identifier.ends_with(">") { // The argument is a Discord mention. let user_id: u64 = identifier[2..identifier.len() - 1] .to_string() .parse() .map_err(|_| ArgumentParsingError::InvalidArgument)?; user = get_user_by_id(user_id, &ctx.rt.state.db).await?; } else if identifier.chars().all(|c| c.is_ascii_digit()) { // The argument is a user ID. let user_id: u64 = identifier .parse() .map_err(|_| ArgumentParsingError::InvalidArgument)?; match get_user_by_id(user_id, &ctx.rt.state.db).await { Ok(result) => user = result, Err(error) => { if matches!(error, ArgumentParsingError::InvalidArgument) { let message = Message::get_by_propagation_message_id(user_id, &ctx.rt.state.db) .await?; if let Some(message) = message { user = get_user_by_id(message.author_id, &ctx.rt.state.db).await?; } else { return Err(ArgumentParsingError::InvalidArgument); } } else { return Err(error); } } } } else { // The argument is (probably) a username. let Some(db_user) = User::get_by_name(identifier.to_string(), &ctx.rt.state.db).await? else { // Check in the message's reply, and if not use the command executor. user = user_from_reply_or_command_executor(ctx).await?; return Ok((UserArgument(user), args)); }; user = db_user; } // Notice that `args` is shadowed above, so this line can't be removed (or else the // full args will be passed and not just the remaining ones as it should be). Ok((UserArgument(user), args)) } else { // Check in the message's reply, and if not use the command executor. let user = user_from_reply_or_command_executor(ctx).await?; Ok((UserArgument(user), args)) } } } impl Deref for UserArgument { type Target = User; fn deref(&self) -> &Self::Target { &self.0 } } impl DerefMut for UserArgument { fn deref_mut(&mut self) -> &mut Self::Target { &mut self.0 } } #[derive(Debug, Clone, Copy)] pub struct DurationArgument(pub Duration); impl FromArguments for DurationArgument { async fn from_args<'a>( mut args: &'a str, _ctx: &'a CommandContext<'_, State, Extras>, ) -> Result<(Self, &'a str), ArgumentParsingError> { let mut i = 1; let mut finding = None; while let Some((substring, remaining_args)) = split_at_nth_from_end(args, " ", i) { if let Ok(duration) = humantime::parse_duration(substring) { finding = Some(DurationArgument(duration)); args = remaining_args; break; } i += 1; } if let Some(duration) = finding { Ok((duration, args)) } else { Err(ArgumentParsingError::InvalidArgument) } } } impl Deref for DurationArgument { type Target = Duration; fn deref(&self) -> &Self::Target { &self.0 } } impl DerefMut for DurationArgument { fn deref_mut(&mut self) -> &mut Self::Target { &mut self.0 } } #[derive(Debug, Clone)] pub struct ChannelArgument(pub Channel); impl FromArguments for ChannelArgument { async fn from_args<'a>( args: &'a str, ctx: &'a CommandContext<'_, State, Extras>, ) -> Result<(Self, &'a str), ArgumentParsingError> { let (identifier, args) = NeedyString::from_args(args, ctx).await?; let identifier = identifier.0.trim(); if identifier == "." || identifier == "this" || identifier == "here" { let channel = Channel::get_by_id(ctx.event.channel_id.get(), &ctx.rt.state.db).await?; if let Some(channel) = channel { return Ok((Self(channel), args)); } } else if identifier.starts_with("<#") && identifier.ends_with(">") { // The argument is a Discord mention. let channel_id: u64 = identifier[2..identifier.len() - 1] .to_string() .parse() .map_err(|_| ArgumentParsingError::InvalidArgument)?; let channel = Channel::get_by_id(channel_id, &ctx.rt.state.db).await?; if let Some(channel) = channel { return Ok((Self(channel), args)); } } else { let channel_id: u64 = identifier .parse() .map_err(|_| ArgumentParsingError::InvalidArgument)?; let channel = Channel::get_by_id(channel_id, &ctx.rt.state.db).await?; if let Some(channel) = channel { return Ok((Self(channel), args)); } } Err(ArgumentParsingError::MissingArgument) } } impl Deref for ChannelArgument { type Target = Channel; fn deref(&self) -> &Self::Target { &self.0 } } impl DerefMut for ChannelArgument { fn deref_mut(&mut self) -> &mut Self::Target { &mut self.0 } }