1use crate::bot::Bot;
6use crate::ctx::{Ctx, StateMap};
7use crate::event_trigger::EventKind;
8use crate::extract::{Extracted, FromContext, Reject, Reply};
9use async_trait::async_trait;
10use nagisa_types::event::Event;
11use nagisa_types::message::MessageExt;
12use nagisa_types::prelude::*;
13use std::any::TypeId;
14use std::collections::HashSet;
15use std::future::Future;
16use std::sync::atomic::{AtomicU64, Ordering};
17use std::sync::{Arc, Mutex};
18use std::time::Duration;
19use tokio::sync::mpsc;
20use tokio::sync::mpsc::error::TrySendError;
21use tokio::time::Instant;
22
23pub type EventFilter = Arc<dyn Fn(&Event) -> bool + Send + Sync>;
26
27#[derive(Clone)]
30pub struct Selector {
31 pub kind: EventKind,
32 pub peer: Option<Peer>,
33 pub user: Option<Uin>,
34 pub filter: Option<EventFilter>,
35}
36
37impl Default for Selector {
38 fn default() -> Self {
39 Selector { kind: EventKind::Message, peer: None, user: None, filter: None }
40 }
41}
42
43impl Selector {
44 pub fn matches(&self, event: &Event) -> bool {
46 if EventKind::of(event) != Some(self.kind) {
47 return false;
48 }
49 if let Some(p) = self.peer {
50 if event.peer() != Some(p) {
51 return false;
52 }
53 }
54 if let Some(u) = self.user {
55 if event.sender() != Some(u) {
56 return false;
57 }
58 }
59 if let Some(f) = &self.filter {
60 if !f(event) {
61 return false;
62 }
63 }
64 true
65 }
66}
67
68#[derive(Clone)]
70pub struct Scope(Selector);
71
72impl Scope {
73 pub fn peer(peer: Peer) -> Scope {
79 Scope(Selector { peer: Some(peer), ..Selector::default() })
80 }
81 pub fn user(user: Uin) -> Scope {
83 Scope(Selector { user: Some(user), ..Selector::default() })
84 }
85 pub fn from(peer: Peer, user: Uin) -> Scope {
87 Scope(Selector { peer: Some(peer), user: Some(user), ..Selector::default() })
88 }
89 pub fn into_selector(self) -> Selector {
91 self.0
92 }
93}
94
95#[derive(Clone, Copy, Debug, PartialEq, Eq)]
97pub struct Delivery {
98 pub delivered: bool,
99 pub block: bool,
100}
101
102pub enum WaitFlow<T> {
109 Continue,
111 Done(T),
113 Cancel,
115}
116
117pub enum Replied<T> {
121 Got(T),
123 Cancelled,
125 TimedOut,
127}
128
129struct WaiterEntry {
131 id: u64,
132 depth: u32,
133 selector: Selector,
134 block: bool,
135 tx: mpsc::Sender<Arc<Event>>,
136}
137
138pub struct WaiterStore {
141 waiters: Mutex<Vec<WaiterEntry>>,
142 next_id: AtomicU64,
143}
144
145impl Default for WaiterStore {
146 fn default() -> Self {
147 Self { waiters: Mutex::new(Vec::new()), next_id: AtomicU64::new(1) }
148 }
149}
150
151impl WaiterStore {
152 pub fn new() -> Self {
153 Self::default()
154 }
155
156 fn lock(&self) -> std::sync::MutexGuard<'_, Vec<WaiterEntry>> {
157 self.waiters.lock().unwrap_or_else(|e| e.into_inner())
158 }
159
160 pub fn next_id(&self) -> u64 {
162 self.next_id.fetch_add(1, Ordering::Relaxed)
163 }
164
165 pub fn register(&self, depth: u32, selector: Selector, block: bool, tx: mpsc::Sender<Arc<Event>>) -> u64 {
167 let id = self.next_id();
168 self.register_with_id(id, depth, selector, block, tx);
169 id
170 }
171
172 pub fn register_with_id(&self, id: u64, depth: u32, selector: Selector, block: bool, tx: mpsc::Sender<Arc<Event>>) {
174 self.lock().push(WaiterEntry { id, depth, selector, block, tx });
175 }
176
177 pub fn remove(&self, id: u64) {
179 self.lock().retain(|e| e.id != id);
180 }
181
182 pub fn len(&self) -> usize {
184 self.lock().len()
185 }
186
187 pub fn is_empty(&self) -> bool {
189 self.lock().is_empty()
190 }
191
192 pub fn try_deliver(&self, event: &Arc<Event>) -> Delivery {
194 loop {
195 let chosen = {
197 let guard = self.lock();
198 guard
199 .iter()
200 .filter(|e| e.selector.matches(event))
201 .max_by_key(|e| (e.depth, e.id))
202 .map(|e| (e.id, e.block, e.tx.clone()))
203 };
204 let Some((id, block, tx)) = chosen else {
205 return Delivery { delivered: false, block: false };
206 };
207 match tx.try_send(Arc::clone(event)) {
208 Ok(()) => return Delivery { delivered: true, block },
209 Err(TrySendError::Closed(_)) => {
211 self.remove(id);
212 continue;
213 }
214 Err(TrySendError::Full(_)) => return Delivery { delivered: false, block: false },
217 }
218 }
219 }
220}
221
222#[derive(Clone, PartialEq, Eq, Hash, Debug)]
226struct FlightKey {
227 kind: EventKind,
228 peer: Option<Peer>,
229 user: Option<Uin>,
230}
231
232impl FlightKey {
233 fn from_selector(sel: &Selector) -> Self {
234 FlightKey { kind: sel.kind, peer: sel.peer, user: sel.user }
235 }
236}
237
238pub struct FlightStore {
243 held: Mutex<HashSet<FlightKey>>,
244}
245
246impl Default for FlightStore {
247 fn default() -> Self {
248 Self { held: Mutex::new(HashSet::new()) }
249 }
250}
251
252impl FlightStore {
253 pub fn new() -> Self {
255 Self::default()
256 }
257
258 fn lock(&self) -> std::sync::MutexGuard<'_, HashSet<FlightKey>> {
259 self.held.lock().unwrap_or_else(|e| e.into_inner())
260 }
261
262 pub fn len(&self) -> usize {
264 self.lock().len()
265 }
266
267 pub fn is_empty(&self) -> bool {
269 self.lock().is_empty()
270 }
271
272 fn acquire(store: Arc<FlightStore>, selector: Selector) -> Option<FlightGuard> {
275 let key = FlightKey::from_selector(&selector);
276 {
277 let mut held = store.lock();
278 if !held.insert(key.clone()) {
279 return None; }
281 }
282 Some(FlightGuard { store, key })
283 }
284}
285
286pub struct FlightGuard {
291 store: Arc<FlightStore>,
292 key: FlightKey,
293}
294
295impl Drop for FlightGuard {
296 fn drop(&mut self) {
297 self.store.lock().remove(&self.key);
298 }
299}
300
301#[derive(Clone, Copy, Debug)]
303pub struct WaiterDepth(pub u32);
304
305pub struct WaiterBuilder {
308 bot: Bot,
309 state: Arc<StateMap>,
310 store: Arc<WaiterStore>,
311 depth: u32,
312 selector: Selector,
313 block: bool,
314 starter_peer: Option<Peer>,
318 starter_user: Option<Uin>,
319}
320
321impl WaiterBuilder {
322 pub fn on(mut self, kind: EventKind) -> Self {
324 self.selector.kind = kind;
325 self
326 }
327 pub fn scope(mut self, scope: Scope) -> Self {
329 self.selector = scope.into_selector();
330 self
331 }
332 pub fn peer(mut self, peer: Peer) -> Self {
334 self.selector.peer = Some(peer);
335 self
336 }
337 pub fn user(mut self, user: Uin) -> Self {
339 self.selector.user = Some(user);
340 self
341 }
342 pub fn from(mut self, peer: Peer, user: Uin) -> Self {
344 self.selector.peer = Some(peer);
345 self.selector.user = Some(user);
346 self
347 }
348 pub fn from_starter(mut self) -> Self {
353 self.selector.peer = self.starter_peer;
354 self.selector.user = self.starter_user;
355 self
356 }
357 pub fn filter<F>(mut self, f: F) -> Self
359 where
360 F: Fn(&Event) -> bool + Send + Sync + 'static,
361 {
362 self.selector.filter = Some(Arc::new(f));
363 self
364 }
365 pub fn block(mut self, block: bool) -> Self {
367 self.block = block;
368 self
369 }
370 pub fn build(self) -> Waiter {
372 let (tx, rx) = mpsc::channel(16);
373 let id = self.store.next_id();
374 self.store.register_with_id(id, self.depth, self.selector, self.block, tx);
375 Waiter {
376 id,
377 rx: tokio::sync::Mutex::new(rx),
378 store: self.store,
379 bot: self.bot,
380 state: self.state,
381 depth: self.depth,
382 }
383 }
384}
385
386pub struct Waiter {
388 id: u64,
389 rx: tokio::sync::Mutex<mpsc::Receiver<Arc<Event>>>,
390 store: Arc<WaiterStore>,
391 bot: Bot,
392 state: Arc<StateMap>,
393 depth: u32,
394}
395
396impl Drop for Waiter {
397 fn drop(&mut self) {
398 self.store.remove(self.id);
399 }
400}
401
402impl Waiter {
403 pub async fn recv_event(&self, timeout: Duration) -> Option<Ctx> {
406 let ev = self.recv_raw(timeout).await?;
407 let ctx = Ctx::new(ev, self.bot.clone(), Arc::clone(&self.state));
408 ctx.insert_ext(WaiterDepth(self.depth));
409 Some(ctx)
410 }
411
412 pub async fn recv<T: FromContext>(&self, timeout: Duration) -> Option<T> {
415 Some(self.recv_session::<T>(timeout).await?.0)
416 }
417
418 pub async fn recv_session<T: FromContext>(&self, timeout: Duration) -> Option<(T, Session)> {
422 let deadline = tokio::time::Instant::now() + timeout;
423 loop {
424 let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
425 if remaining.is_zero() {
426 return None;
427 }
428 let ctx = self.recv_event(remaining).await?;
429 match T::from_context(&ctx).await {
430 Ok(v) => match Session::from_context(&ctx).await {
433 Ok(session) => return Some((v, session)),
434 Err(_) => return None,
435 },
436 Err(Reject::Skip) => continue,
437 Err(Reject::Error(_)) => return None,
438 }
439 }
440 }
441
442 pub async fn recv_with<T, F, Fut>(&self, timeout: Duration, mut f: F) -> Option<T>
473 where
474 F: FnMut(Ctx) -> Fut,
475 Fut: Future<Output = WaitFlow<T>>,
476 {
477 let deadline = Instant::now() + timeout;
478 loop {
479 let remaining = deadline.saturating_duration_since(Instant::now());
480 if remaining.is_zero() {
481 return None;
482 }
483 let ctx = self.recv_event(remaining).await?;
484 match f(ctx).await {
485 WaitFlow::Continue => continue,
486 WaitFlow::Done(v) => return Some(v),
487 WaitFlow::Cancel => return None,
488 }
489 }
490 }
491
492 pub async fn recv_with_sync<T, F>(&self, timeout: Duration, mut f: F) -> Option<T>
497 where
498 F: FnMut(Ctx) -> WaitFlow<T>,
499 {
500 let deadline = Instant::now() + timeout;
501 loop {
502 let remaining = deadline.saturating_duration_since(Instant::now());
503 if remaining.is_zero() {
504 return None;
505 }
506 let ctx = self.recv_event(remaining).await?;
507 match f(ctx) {
508 WaitFlow::Continue => continue,
509 WaitFlow::Done(v) => return Some(v),
510 WaitFlow::Cancel => return None,
511 }
512 }
513 }
514
515 pub async fn confirm(&self, timeout: Duration, who: Uin, yes: &str, no: &str) -> Option<bool> {
521 let yes = yes.trim().to_lowercase();
522 let no = no.trim().to_lowercase();
523 self.recv_with(timeout, |ctx| {
524 let yes = yes.clone();
525 let no = no.clone();
526 async move {
527 let Some(m) = ctx.message() else { return WaitFlow::Continue };
528 if m.sender != who {
529 return WaitFlow::Continue;
530 }
531 let text = m.content.extract_text();
532 let t = text.trim().to_lowercase();
533 if t == yes {
534 WaitFlow::Done(true)
535 } else if t == no {
536 WaitFlow::Done(false)
537 } else {
538 if let Ok(reply) = Reply::from_context(&ctx).await {
540 let _ = reply.text(format!("请回复「{yes}」或「{no}」")).await;
541 }
542 WaitFlow::Continue
543 }
544 }
545 })
546 .await
547 }
548
549 pub async fn recv_text(&self, timeout: Duration, cancel: impl Fn(&str) -> bool) -> Replied<String> {
553 self.recv_with_sync(timeout, |ctx| {
554 let Some(m) = ctx.message() else { return WaitFlow::Continue };
555 let text = m.content.extract_text();
556 let trimmed = text.trim();
557 if trimmed.is_empty() {
558 return WaitFlow::Continue;
559 }
560 if cancel(trimmed) {
561 return WaitFlow::Done(Replied::Cancelled);
562 }
563 WaitFlow::Done(Replied::Got(text))
564 })
565 .await
566 .unwrap_or(Replied::TimedOut)
567 }
568
569 pub async fn recv_parse<T, F, C>(&self, timeout: Duration, cancel: C, parse: F) -> Replied<T>
578 where
579 F: Fn(&str) -> std::result::Result<T, String>,
580 C: Fn(&str) -> bool,
581 {
582 enum Pre<U> {
585 Skip,
586 Cancel,
587 Done(U),
588 Reprompt(String),
589 }
590 self.recv_with(timeout, move |ctx| {
591 let pre = match ctx.message() {
592 None => Pre::Skip,
593 Some(m) => {
594 let text = m.content.extract_text();
595 let trimmed = text.trim();
596 if trimmed.is_empty() {
597 Pre::Skip
598 } else if cancel(trimmed) {
599 Pre::Cancel
600 } else {
601 match parse(trimmed) {
602 Ok(v) => Pre::Done(v),
603 Err(hint) => Pre::Reprompt(hint),
604 }
605 }
606 }
607 };
608 async move {
609 match pre {
610 Pre::Skip => WaitFlow::Continue,
611 Pre::Cancel => WaitFlow::Done(Replied::Cancelled),
612 Pre::Done(v) => WaitFlow::Done(Replied::Got(v)),
613 Pre::Reprompt(hint) => {
614 if let Ok(reply) = Reply::from_context(&ctx).await {
615 let _ = reply.text(hint).await;
616 }
617 WaitFlow::Continue
618 }
619 }
620 }
621 })
622 .await
623 .unwrap_or(Replied::TimedOut)
624 }
625
626 async fn rx_guard(&self) -> tokio::sync::MutexGuard<'_, mpsc::Receiver<Arc<Event>>> {
627 self.rx.lock().await
628 }
629
630 async fn recv_raw(&self, timeout: Duration) -> Option<Arc<Event>> {
631 let mut rx = self.rx_guard().await;
632 match tokio::time::timeout(timeout, rx.recv()).await {
633 Ok(Some(ev)) => Some(ev),
634 _ => None,
635 }
636 }
637}
638
639pub struct Session {
642 bot: Bot,
643 state: Arc<StateMap>,
644 store: Arc<WaiterStore>,
645 flight: Arc<FlightStore>,
646 peer: Option<Peer>,
647 user: Option<Uin>,
648 depth: u32,
649}
650
651impl Session {
652 pub fn peer(&self) -> Option<Peer> {
654 self.peer
655 }
656
657 pub fn user(&self) -> Option<Uin> {
659 self.user
660 }
661
662 pub fn waiter(&self) -> WaiterBuilder {
664 WaiterBuilder {
665 bot: self.bot.clone(),
666 state: Arc::clone(&self.state),
667 store: Arc::clone(&self.store),
668 depth: self.depth + 1,
669 selector: Selector::default(),
670 block: true,
671 starter_peer: self.peer,
672 starter_user: self.user,
673 }
674 }
675
676 pub fn single_flight(&self, scope: Scope) -> Option<FlightGuard> {
683 FlightStore::acquire(Arc::clone(&self.flight), scope.into_selector())
684 }
685
686 pub fn single_flight_user(&self) -> Option<FlightGuard> {
691 self.user.and_then(|u| self.single_flight(Scope::user(u)))
692 }
693}
694
695#[async_trait]
696impl FromContext for Session {
697 async fn from_context(ctx: &Ctx) -> Extracted<Self> {
698 let store = ctx
699 .state()
700 .get(&TypeId::of::<WaiterStore>())
701 .and_then(|a| Arc::clone(a).downcast::<WaiterStore>().ok())
702 .ok_or_else(|| {
703 Reject::Error(Error::action_kind(
704 ActionErrorKind::Other,
705 "WaiterStore not registered (use App::new or register it)",
706 ))
707 })?;
708 let flight = ctx
711 .state()
712 .get(&TypeId::of::<FlightStore>())
713 .and_then(|a| Arc::clone(a).downcast::<FlightStore>().ok())
714 .unwrap_or_else(|| Arc::new(FlightStore::new()));
715 let depth = ctx.get_ext::<WaiterDepth>().map(|d| d.0).unwrap_or(0);
716 Ok(Session {
717 bot: ctx.bot().clone(),
718 state: Arc::clone(ctx.state()),
719 store,
720 flight,
721 peer: ctx.event().peer(),
722 user: ctx.event().sender(),
723 depth,
724 })
725 }
726}