payphone/src/groups/moderation/arguments.rs

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<User, ArgumentParser>`] - 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<User, ArgumentParsingError> {
    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<User, ArgumentParsingError>`] - 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<User, ArgumentParsingError> {
    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<State, Extras> 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<State, Extras> 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<State, Extras> 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
    }
}