use crate::bot::Bot;
use crate::ctx::{Ctx, StateMap};
use crate::event_trigger::EventKind;
use crate::extract::{Extracted, FromContext, Reject, Reply};
use async_trait::async_trait;
use nagisa_types::event::Event;
use nagisa_types::prelude::*;
use nagisa_types::message::MessageExt;
use std::any::TypeId;
use std::collections::HashSet;
use std::future::Future;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::mpsc;
use tokio::sync::mpsc::error::TrySendError;
use tokio::time::Instant;
pub type EventFilter = Arc<dyn Fn(&Event) -> bool + Send + Sync>;
#[derive(Clone)]
pub struct Selector {
pub kind: EventKind,
pub peer: Option<Peer>,
pub user: Option<Uin>,
pub filter: Option<EventFilter>,
}
impl Default for Selector {
fn default() -> Self {
Selector { kind: EventKind::Message, peer: None, user: None, filter: None }
}
}
impl Selector {
pub fn matches(&self, event: &Event) -> bool {
if EventKind::of(event) != Some(self.kind) {
return false;
}
if let Some(p) = self.peer {
if event.peer() != Some(p) {
return false;
}
}
if let Some(u) = self.user {
if event.sender() != Some(u) {
return false;
}
}
if let Some(f) = &self.filter {
if !f(event) {
return false;
}
}
true
}
}
#[derive(Clone)]
pub struct Scope(Selector);
impl Scope {
pub fn peer(peer: Peer) -> Scope {
Scope(Selector { peer: Some(peer), ..Selector::default() })
}
pub fn user(user: Uin) -> Scope {
Scope(Selector { user: Some(user), ..Selector::default() })
}
pub fn from(peer: Peer, user: Uin) -> Scope {
Scope(Selector { peer: Some(peer), user: Some(user), ..Selector::default() })
}
pub fn into_selector(self) -> Selector {
self.0
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Delivery {
pub delivered: bool,
pub block: bool,
}
pub enum WaitFlow<T> {
Continue,
Done(T),
Cancel,
}
struct WaiterEntry {
id: u64,
depth: u32,
selector: Selector,
block: bool,
tx: mpsc::Sender<Arc<Event>>,
}
pub struct WaiterStore {
waiters: Mutex<Vec<WaiterEntry>>,
next_id: AtomicU64,
}
impl Default for WaiterStore {
fn default() -> Self {
Self { waiters: Mutex::new(Vec::new()), next_id: AtomicU64::new(1) }
}
}
impl WaiterStore {
pub fn new() -> Self {
Self::default()
}
fn lock(&self) -> std::sync::MutexGuard<'_, Vec<WaiterEntry>> {
self.waiters.lock().unwrap_or_else(|e| e.into_inner())
}
pub fn next_id(&self) -> u64 {
self.next_id.fetch_add(1, Ordering::Relaxed)
}
pub fn register(
&self,
depth: u32,
selector: Selector,
block: bool,
tx: mpsc::Sender<Arc<Event>>,
) -> u64 {
let id = self.next_id();
self.register_with_id(id, depth, selector, block, tx);
id
}
pub fn register_with_id(
&self,
id: u64,
depth: u32,
selector: Selector,
block: bool,
tx: mpsc::Sender<Arc<Event>>,
) {
self.lock().push(WaiterEntry { id, depth, selector, block, tx });
}
pub fn remove(&self, id: u64) {
self.lock().retain(|e| e.id != id);
}
pub fn len(&self) -> usize {
self.lock().len()
}
pub fn is_empty(&self) -> bool {
self.lock().is_empty()
}
pub fn try_deliver(&self, event: &Arc<Event>) -> Delivery {
loop {
let chosen = {
let guard = self.lock();
guard
.iter()
.filter(|e| e.selector.matches(event))
.max_by_key(|e| (e.depth, e.id))
.map(|e| (e.id, e.block, e.tx.clone()))
};
let Some((id, block, tx)) = chosen else {
return Delivery { delivered: false, block: false };
};
match tx.try_send(Arc::clone(event)) {
Ok(()) => return Delivery { delivered: true, block },
Err(TrySendError::Closed(_)) => {
self.remove(id);
continue;
}
Err(TrySendError::Full(_)) => return Delivery { delivered: false, block: false },
}
}
}
}
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
struct FlightKey {
kind: EventKind,
peer: Option<Peer>,
user: Option<Uin>,
}
impl FlightKey {
fn from_selector(sel: &Selector) -> Self {
FlightKey { kind: sel.kind, peer: sel.peer, user: sel.user }
}
}
pub struct FlightStore {
held: Mutex<HashSet<FlightKey>>,
}
impl Default for FlightStore {
fn default() -> Self {
Self { held: Mutex::new(HashSet::new()) }
}
}
impl FlightStore {
pub fn new() -> Self {
Self::default()
}
fn lock(&self) -> std::sync::MutexGuard<'_, HashSet<FlightKey>> {
self.held.lock().unwrap_or_else(|e| e.into_inner())
}
pub fn len(&self) -> usize {
self.lock().len()
}
pub fn is_empty(&self) -> bool {
self.lock().is_empty()
}
fn acquire(store: Arc<FlightStore>, selector: Selector) -> Option<FlightGuard> {
let key = FlightKey::from_selector(&selector);
{
let mut held = store.lock();
if !held.insert(key.clone()) {
return None; }
}
Some(FlightGuard { store, key })
}
}
pub struct FlightGuard {
store: Arc<FlightStore>,
key: FlightKey,
}
impl Drop for FlightGuard {
fn drop(&mut self) {
self.store.lock().remove(&self.key);
}
}
#[derive(Clone, Copy, Debug)]
pub struct WaiterDepth(pub u32);
pub struct WaiterBuilder {
bot: Bot,
state: Arc<StateMap>,
store: Arc<WaiterStore>,
depth: u32,
selector: Selector,
block: bool,
starter_peer: Option<Peer>,
starter_user: Option<Uin>,
}
impl WaiterBuilder {
pub fn on(mut self, kind: EventKind) -> Self {
self.selector.kind = kind;
self
}
pub fn scope(mut self, scope: Scope) -> Self {
self.selector = scope.into_selector();
self
}
pub fn peer(mut self, peer: Peer) -> Self {
self.selector.peer = Some(peer);
self
}
pub fn user(mut self, user: Uin) -> Self {
self.selector.user = Some(user);
self
}
pub fn from(mut self, peer: Peer, user: Uin) -> Self {
self.selector.peer = Some(peer);
self.selector.user = Some(user);
self
}
pub fn from_starter(mut self) -> Self {
self.selector.peer = self.starter_peer;
self.selector.user = self.starter_user;
self
}
pub fn filter<F>(mut self, f: F) -> Self
where
F: Fn(&Event) -> bool + Send + Sync + 'static,
{
self.selector.filter = Some(Arc::new(f));
self
}
pub fn block(mut self, block: bool) -> Self {
self.block = block;
self
}
pub fn build(self) -> Waiter {
let (tx, rx) = mpsc::channel(16);
let id = self.store.next_id();
self.store.register_with_id(id, self.depth, self.selector, self.block, tx);
Waiter {
id,
rx: tokio::sync::Mutex::new(rx),
store: self.store,
bot: self.bot,
state: self.state,
depth: self.depth,
}
}
}
pub struct Waiter {
id: u64,
rx: tokio::sync::Mutex<mpsc::Receiver<Arc<Event>>>,
store: Arc<WaiterStore>,
bot: Bot,
state: Arc<StateMap>,
depth: u32,
}
impl Drop for Waiter {
fn drop(&mut self) {
self.store.remove(self.id);
}
}
impl Waiter {
pub async fn recv_event(&self, timeout: Duration) -> Option<Ctx> {
let ev = self.recv_raw(timeout).await?;
let ctx = Ctx::new(ev, self.bot.clone(), Arc::clone(&self.state));
ctx.insert_ext(WaiterDepth(self.depth));
Some(ctx)
}
pub async fn recv<T: FromContext>(&self, timeout: Duration) -> Option<T> {
Some(self.recv_session::<T>(timeout).await?.0)
}
pub async fn recv_session<T: FromContext>(&self, timeout: Duration) -> Option<(T, Session)> {
let deadline = tokio::time::Instant::now() + timeout;
loop {
let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
if remaining.is_zero() {
return None;
}
let ctx = self.recv_event(remaining).await?;
match T::from_context(&ctx).await {
Ok(v) => match Session::from_context(&ctx).await {
Ok(session) => return Some((v, session)),
Err(_) => return None,
},
Err(Reject::Skip) => continue,
Err(Reject::Error(_)) => return None,
}
}
}
pub async fn recv_with<T, F, Fut>(&self, timeout: Duration, mut f: F) -> Option<T>
where
F: FnMut(Ctx) -> Fut,
Fut: Future<Output = WaitFlow<T>>,
{
let deadline = Instant::now() + timeout;
loop {
let remaining = deadline.saturating_duration_since(Instant::now());
if remaining.is_zero() {
return None;
}
let ctx = self.recv_event(remaining).await?;
match f(ctx).await {
WaitFlow::Continue => continue,
WaitFlow::Done(v) => return Some(v),
WaitFlow::Cancel => return None,
}
}
}
pub async fn recv_with_sync<T, F>(&self, timeout: Duration, mut f: F) -> Option<T>
where
F: FnMut(Ctx) -> WaitFlow<T>,
{
let deadline = Instant::now() + timeout;
loop {
let remaining = deadline.saturating_duration_since(Instant::now());
if remaining.is_zero() {
return None;
}
let ctx = self.recv_event(remaining).await?;
match f(ctx) {
WaitFlow::Continue => continue,
WaitFlow::Done(v) => return Some(v),
WaitFlow::Cancel => return None,
}
}
}
pub async fn confirm(
&self,
timeout: Duration,
who: Uin,
yes: &str,
no: &str,
) -> Option<bool> {
let yes = yes.trim().to_lowercase();
let no = no.trim().to_lowercase();
self.recv_with(timeout, |ctx| {
let yes = yes.clone();
let no = no.clone();
async move {
let Some(m) = ctx.message() else { return WaitFlow::Continue };
if m.sender != who {
return WaitFlow::Continue;
}
let text = m.content.extract_text();
let t = text.trim().to_lowercase();
if t == yes {
WaitFlow::Done(true)
} else if t == no {
WaitFlow::Done(false)
} else {
if let Ok(reply) = Reply::from_context(&ctx).await {
let _ = reply.text(format!("请回复「{yes}」或「{no}」")).await;
}
WaitFlow::Continue
}
}
})
.await
}
pub async fn recv_text(&self, timeout: Duration, cancel: &str) -> Option<String> {
let cancel = cancel.trim().to_lowercase();
self.recv_with_sync(timeout, |ctx| {
let Some(m) = ctx.message() else { return WaitFlow::Continue };
let text = m.content.extract_text();
let trimmed = text.trim();
if trimmed.is_empty() {
return WaitFlow::Continue;
}
if trimmed.to_lowercase() == cancel {
return WaitFlow::Cancel;
}
WaitFlow::Done(text)
})
.await
}
pub async fn recv_parse<T, F>(&self, timeout: Duration, cancel: &str, parse: F) -> Option<T>
where
F: Fn(&str) -> std::result::Result<T, String>,
{
enum Pre<U> {
Skip,
Cancel,
Done(U),
Reprompt(String),
}
let cancel = cancel.trim().to_lowercase();
self.recv_with(timeout, move |ctx| {
let pre = match ctx.message() {
None => Pre::Skip,
Some(m) => {
let text = m.content.extract_text();
let trimmed = text.trim();
if trimmed.is_empty() {
Pre::Skip
} else if trimmed.to_lowercase() == cancel {
Pre::Cancel
} else {
match parse(trimmed) {
Ok(v) => Pre::Done(v),
Err(hint) => Pre::Reprompt(hint),
}
}
}
};
async move {
match pre {
Pre::Skip => WaitFlow::Continue,
Pre::Cancel => WaitFlow::Cancel,
Pre::Done(v) => WaitFlow::Done(v),
Pre::Reprompt(hint) => {
if let Ok(reply) = Reply::from_context(&ctx).await {
let _ = reply.text(hint).await;
}
WaitFlow::Continue
}
}
}
})
.await
}
async fn rx_guard(&self) -> tokio::sync::MutexGuard<'_, mpsc::Receiver<Arc<Event>>> {
self.rx.lock().await
}
async fn recv_raw(&self, timeout: Duration) -> Option<Arc<Event>> {
let mut rx = self.rx_guard().await;
match tokio::time::timeout(timeout, rx.recv()).await {
Ok(Some(ev)) => Some(ev),
_ => None,
}
}
}
pub struct Session {
bot: Bot,
state: Arc<StateMap>,
store: Arc<WaiterStore>,
flight: Arc<FlightStore>,
peer: Option<Peer>,
user: Option<Uin>,
depth: u32,
}
impl Session {
pub fn peer(&self) -> Option<Peer> {
self.peer
}
pub fn user(&self) -> Option<Uin> {
self.user
}
pub fn waiter(&self) -> WaiterBuilder {
WaiterBuilder {
bot: self.bot.clone(),
state: Arc::clone(&self.state),
store: Arc::clone(&self.store),
depth: self.depth + 1,
selector: Selector::default(),
block: true,
starter_peer: self.peer,
starter_user: self.user,
}
}
pub fn single_flight(&self, scope: Scope) -> Option<FlightGuard> {
FlightStore::acquire(Arc::clone(&self.flight), scope.into_selector())
}
pub fn single_flight_user(&self) -> Option<FlightGuard> {
self.user.and_then(|u| self.single_flight(Scope::user(u)))
}
}
#[async_trait]
impl FromContext for Session {
async fn from_context(ctx: &Ctx) -> Extracted<Self> {
let store = ctx
.state()
.get(&TypeId::of::<WaiterStore>())
.and_then(|a| Arc::clone(a).downcast::<WaiterStore>().ok())
.ok_or_else(|| {
Reject::Error(Error::action_kind(
ActionErrorKind::Other,
"WaiterStore not registered (use App::new or register it)",
))
})?;
let flight = ctx
.state()
.get(&TypeId::of::<FlightStore>())
.and_then(|a| Arc::clone(a).downcast::<FlightStore>().ok())
.unwrap_or_else(|| Arc::new(FlightStore::new()));
let depth = ctx.get_ext::<WaiterDepth>().map(|d| d.0).unwrap_or(0);
Ok(Session {
bot: ctx.bot().clone(),
state: Arc::clone(ctx.state()),
store,
flight,
peer: ctx.event().peer(),
user: ctx.event().sender(),
depth,
})
}
}