use crate::ctx::Ctx;
use crate::extract::{Extracted, FromContext, Reject, State};
use crate::rule::{Check, Rule};
use async_trait::async_trait;
use nagisa_types::id::{Peer, Uin};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CooldownScope {
UserGlobal,
User,
Peer,
Global,
}
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct TriggerId {
plugin: String,
trigger: String,
}
impl TriggerId {
pub fn of(plugin: impl Into<String>, trigger: impl Into<String>) -> Self {
Self { plugin: plugin.into(), trigger: trigger.into() }
}
}
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub enum CdKey {
UserGlobal(Uin),
User(Uin, TriggerId),
Peer(Peer),
Global,
Custom(String),
}
impl From<&str> for CdKey {
fn from(s: &str) -> Self {
CdKey::Custom(s.to_string())
}
}
impl From<String> for CdKey {
fn from(s: String) -> Self {
CdKey::Custom(s)
}
}
#[derive(Clone, Copy, Debug)]
struct Window {
count: u32,
first: Instant,
}
#[derive(Default)]
pub struct CooldownStore {
map: Mutex<HashMap<CdKey, Window>>,
}
impl CooldownStore {
pub fn new() -> Self {
Self::default()
}
fn try_stamp(&self, key: CdKey, interval: Duration, max_exec: u32) -> bool {
let now = Instant::now();
let mut map = self.map.lock().unwrap_or_else(|e| e.into_inner());
match map.get_mut(&key) {
Some(w) if now.duration_since(w.first) < interval => {
if w.count < max_exec {
w.count += 1;
true
} else {
false
}
}
_ => {
map.insert(key, Window { count: 1, first: now });
true
}
}
}
fn ready(&self, key: &CdKey, interval: Duration, max_exec: u32) -> bool {
let now = Instant::now();
let map = self.map.lock().unwrap_or_else(|e| e.into_inner());
match map.get(key) {
Some(w) if now.duration_since(w.first) < interval => w.count < max_exec,
_ => true,
}
}
fn stamp(&self, key: CdKey, interval: Duration) {
let now = Instant::now();
let mut map = self.map.lock().unwrap_or_else(|e| e.into_inner());
match map.get_mut(&key) {
Some(w) if now.duration_since(w.first) < interval => w.count += 1,
_ => {
map.insert(key, Window { count: 1, first: now });
}
}
}
}
#[derive(Clone)]
pub struct Cooldown {
interval: Duration,
scope: CooldownScope,
max_exec: u32,
bypass: Option<Rule>,
}
impl Cooldown {
pub fn new(secs: u64) -> Self {
Self { interval: Duration::from_secs(secs), scope: CooldownScope::UserGlobal, max_exec: 1, bypass: None }
}
pub fn per_peer(mut self) -> Self {
self.scope = CooldownScope::Peer;
self
}
pub fn per_trigger(mut self) -> Self {
self.scope = CooldownScope::User;
self
}
pub fn global(mut self) -> Self {
self.scope = CooldownScope::Global;
self
}
pub fn max_exec(mut self, n: u32) -> Self {
self.max_exec = n.max(1);
self
}
pub fn bypass(mut self, r: Rule) -> Self {
self.bypass = Some(r);
self
}
pub fn into_rule(self, trigger: TriggerId) -> Rule {
Rule::new(CooldownCheck { cd: self, trigger })
}
}
impl From<u64> for Cooldown {
fn from(secs: u64) -> Self {
Cooldown::new(secs)
}
}
fn key_for(ctx: &Ctx, scope: CooldownScope, trigger: &TriggerId) -> Option<CdKey> {
if ctx.message().is_some_and(|m| m.is_self) {
return None;
}
match scope {
CooldownScope::UserGlobal => Some(CdKey::UserGlobal(ctx.message()?.sender)),
CooldownScope::User => Some(CdKey::User(ctx.message()?.sender, trigger.clone())),
CooldownScope::Peer => Some(CdKey::Peer(ctx.message()?.peer)),
CooldownScope::Global => Some(CdKey::Global),
}
}
fn store(ctx: &Ctx) -> Option<Arc<CooldownStore>> {
use std::any::TypeId;
ctx.state().get(&TypeId::of::<CooldownStore>()).and_then(|a| Arc::clone(a).downcast::<CooldownStore>().ok())
}
struct CooldownCheck {
cd: Cooldown,
trigger: TriggerId,
}
#[async_trait]
impl Check for CooldownCheck {
async fn check(&self, ctx: &Ctx) -> bool {
if let Some(b) = &self.cd.bypass {
if b.eval(ctx).await {
return true;
}
}
let Some(key) = key_for(ctx, self.cd.scope, &self.trigger) else {
return true;
};
let Some(store) = store(ctx) else {
return true;
};
store.try_stamp(key, self.cd.interval, self.cd.max_exec)
}
}
pub struct Cd(Arc<CooldownStore>);
#[async_trait]
impl FromContext for Cd {
async fn from_context(ctx: &Ctx) -> Extracted<Self> {
let State(store) = State::<CooldownStore>::from_context(ctx).await?;
Ok(Cd(store))
}
}
impl Cd {
pub fn gate(&self, key: impl Into<CdKey>, dur: Duration) -> Extracted<()> {
if self.0.try_stamp(key.into(), dur, 1) {
Ok(())
} else {
Err(Reject::Skip)
}
}
pub fn ready(&self, key: impl Into<CdKey>, dur: Duration) -> bool {
self.0.ready(&key.into(), dur, 1)
}
pub fn stamp(&self, key: impl Into<CdKey>, dur: Duration) {
self.0.stamp(key.into(), dur);
}
}