flow_bot/base/

extract.rs

use std::{ops::Deref, sync::Arc};

use async_trait::async_trait;

use crate::{
    api::api_ext::ApiExt,
    event::{
        BotEvent, TypedEvent,
        message::{
            GroupSenderInfo, GroupSenderRole, PrivateSenderInfo, SenderSex, TypedMessageInfo,
        },
    },
    message::{self, message_ext::MessageExt, segments::Segment},
};

use super::context::BotContext;

#[async_trait]
/// Extractor trait for extracting information from BotEvent and BotContext.
pub trait FromEvent {
    async fn from_event(context: BotContext, event: BotEvent) -> Option<Self>
    where
        Self: Sized;
}

/// State extractor, extract the state from BotContext.
/// If the required state is not found, the handler will be skipped.
pub struct State<S>(pub Arc<S>);

#[async_trait]
impl<S> FromEvent for State<S>
where
    S: 'static + Send + Sync,
{
    async fn from_event(context: BotContext, _: BotEvent) -> Option<Self> {
        let state = context.state.get::<S>()?;
        Some(Self(state))
    }
}

impl<S> Deref for State<S> {
    type Target = Arc<S>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

/// Extractor for message event.
pub struct MessageBody(pub message::Message);

#[async_trait]
impl FromEvent for MessageBody {
    async fn from_event(_: BotContext, event: BotEvent) -> Option<MessageBody> {
        match event.event {
            TypedEvent::Message(ref msg) => Some(Self(msg.message.clone())),
            _ => None,
        }
    }
}

#[async_trait]
impl FromEvent for GroupSenderRole {
    async fn from_event(_: BotContext, event: BotEvent) -> Option<Self> {
        match event.event {
            TypedEvent::Message(ref msg) => match &msg.info {
                TypedMessageInfo::Group(info) => info.sender.role.clone(),
                _ => None,
            },
            _ => None,
        }
    }
}

#[async_trait]
impl FromEvent for GroupSenderInfo {
    async fn from_event(_: BotContext, event: BotEvent) -> Option<Self> {
        match event.event {
            TypedEvent::Message(ref msg) => match &msg.info {
                TypedMessageInfo::Group(info) => Some(info.sender.clone()),
                _ => None,
            },
            _ => None,
        }
    }
}

#[async_trait]
impl FromEvent for PrivateSenderInfo {
    async fn from_event(_: BotContext, event: BotEvent) -> Option<Self> {
        match event.event {
            TypedEvent::Message(ref msg) => match &msg.info {
                TypedMessageInfo::Private(info) => Some(info.sender.clone()),
                _ => None,
            },
            _ => None,
        }
    }
}

pub struct BasicSenderInfo {
    pub user_id: Option<i64>,
    pub nickname: Option<String>,
    pub sex: Option<SenderSex>,
    pub age: Option<i32>,
}

impl From<PrivateSenderInfo> for BasicSenderInfo {
    fn from(info: PrivateSenderInfo) -> Self {
        Self {
            user_id: info.user_id,
            nickname: info.nickname,
            sex: info.sex,
            age: info.age,
        }
    }
}

impl From<GroupSenderInfo> for BasicSenderInfo {
    fn from(info: GroupSenderInfo) -> Self {
        Self {
            user_id: info.user_id,
            nickname: info.nickname,
            sex: info.sex,
            age: info.age,
        }
    }
}

pub struct Sender(pub BasicSenderInfo);

#[async_trait]
impl FromEvent for Sender {
    async fn from_event(_: BotContext, event: BotEvent) -> Option<Self> {
        match event.event {
            TypedEvent::Message(ref msg) => {
                let info = match &msg.info {
                    TypedMessageInfo::Private(info) => info.sender.clone().into(),
                    TypedMessageInfo::Group(info) => info.sender.clone().into(),
                };
                Some(Self(info))
            }
            _ => None,
        }
    }
}

pub struct At(pub String);

#[async_trait]
impl FromEvent for At {
    async fn from_event(_: BotContext, event: BotEvent) -> Option<Self> {
        if let TypedEvent::Message(ref msg) = event.event {
            msg.message.iter().find_map(|seg| match seg {
                Segment::At(at) => Some(Self(at.qq.clone())),
                _ => None,
            })
        } else {
            None
        }
    }
}

pub struct GroupId(pub i64);

#[async_trait]
impl FromEvent for GroupId {
    async fn from_event(_: BotContext, event: BotEvent) -> Option<Self>
    where
        Self: Sized,
    {
        if let TypedEvent::Message(ref msg) = event.event {
            match &msg.info {
                TypedMessageInfo::Group(info) => Some(Self(info.group_id)),
                _ => None,
            }
        } else {
            None
        }
    }
}

pub struct SenderId(pub i64);

#[async_trait]
impl FromEvent for SenderId {
    async fn from_event(context: BotContext, event: BotEvent) -> Option<Self>
    where
        Self: Sized,
    {
        let sender_info = Sender::from_event(context, event).await?;
        Some(Self(sender_info.0.user_id?))
    }
}

pub struct MatchGroupId<const ID: i64>;

#[async_trait]
impl<const ID: i64> FromEvent for MatchGroupId<ID> {
    async fn from_event(context: BotContext, event: BotEvent) -> Option<Self>
    where
        Self: Sized,
    {
        let group_id = GroupId::from_event(context, event).await?.0;
        if group_id == ID { Some(Self) } else { None }
    }
}

pub struct Reply(pub message::Message);

#[async_trait]
impl FromEvent for Reply {
    async fn from_event(context: BotContext, event: BotEvent) -> Option<Self>
    where
        Self: Sized,
    {
        if let TypedEvent::Message(ref msg) = event.event {
            for segment in msg.message.iter() {
                if let Segment::Reply(reply) = segment {
                    let id = reply.id.parse::<i64>().ok()?;
                    let message = context
                        .get_message(id)
                        .await
                        .map(|msg| msg.message.clone())
                        .ok()?;
                    return Some(Self(message));
                }
            }
        }
        None
    }
}

pub struct Command<const PREFIX: &'static str, Cmd> {
    pub command: Cmd,
}

#[async_trait]
impl<const PREFIX: &'static str, Cmd> FromEvent for Command<PREFIX, Cmd>
where
    Cmd: clap::Parser,
{
    async fn from_event(context: BotContext, event: BotEvent) -> Option<Self>
    where
        Self: Sized,
    {
        let message_body = MessageBody::from_event(context, event).await?;
        let plain_text = message_body.0.extract_if_plain_text()?;
        let trimmed = plain_text.trim();

        if trimmed.split_whitespace().next() != Some(PREFIX) {
            return None;
        }

        let cmd = Cmd::try_parse_from(trimmed.split_whitespace()).ok()?;
        Some(Self { command: cmd })
    }
}

#[async_trait]
impl<T> FromEvent for Option<T>
where
    T: FromEvent,
{
    async fn from_event(context: BotContext, event: BotEvent) -> Option<Self>
    where
        Self: Sized,
    {
        Some(T::from_event(context, event).await)
    }
}

/// A helper macro for matching one of the variants.
/// The macro will generate an enum with the given name and variants.
/// The enum will implement the FromEvent trait, and will try to match the event with the given matchers.
///
/// # Example
/// ```no_run
/// match_one!(MatchOne, A: AMatcher, B: BMatcher);
/// ```
/// The above code will generate an enum like this:
/// ```no_run
/// pub enum MatchOne {
///    A(AMatcher),
///    B(BMatcher),
/// }
/// ```
#[macro_export]
macro_rules! match_one {
    ($name:ident,$($variant:ident : $matcher:ty),*) => {
        pub enum $name {
            $(
                $variant($matcher),
            )*
        }

        #[async_trait::async_trait]
        impl $crate::base::extract::FromEvent for $name {
            async fn from_event(context: $crate::base::context::BotContext, event: $crate::event::BotEvent) -> Option<Self> {
                $(
                    if let Some(matcher) = <$matcher>::from_event(context.clone(), event.clone()).await {
                        return Some(Self::$variant(matcher));
                    }
                )*
                None
            }
        }
    };
}