From d7e28b6d919af127f656557f9552c048fca2a5c5 Mon Sep 17 00:00:00 2001 From: Kishor Date: Thu, 1 Jan 2026 04:13:55 +0530 Subject: [PATCH] added optimization for leaderboard --- Cargo.lock | 343 ++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 3 +- example.env | 7 +- src/api.rs | 355 ++++++++++++++++++++++++++++++++++++++++++ src/commands/level.rs | 246 ++++++++++++++++++++++------- src/listener.rs | 3 + src/main.rs | 77 ++++++++- 7 files changed, 967 insertions(+), 67 deletions(-) create mode 100644 src/api.rs diff --git a/Cargo.lock b/Cargo.lock index 849c25e..e1b6577 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -28,6 +28,189 @@ version = "0.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "366ffbaa4442f4684d91e2cd7c5ea7c4ed8add41959a31447066e279e432b618" +[[package]] +name = "actix-codec" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f7b0a21988c1bf877cf4759ef5ddaac04c1c9fe808c9142ecb78ba97d97a28a" +dependencies = [ + "bitflags 2.10.0", + "bytes", + "futures-core", + "futures-sink", + "memchr", + "pin-project-lite", + "tokio", + "tokio-util", + "tracing", +] + +[[package]] +name = "actix-http" +version = "3.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7926860314cbe2fb5d1f13731e387ab43bd32bca224e82e6e2db85de0a3dba49" +dependencies = [ + "actix-codec", + "actix-rt", + "actix-service", + "actix-utils", + "base64 0.22.1", + "bitflags 2.10.0", + "brotli", + "bytes", + "bytestring", + "derive_more", + "encoding_rs", + "flate2", + "foldhash", + "futures-core", + "h2 0.3.27", + "http 0.2.12", + "httparse", + "httpdate", + "itoa", + "language-tags", + "local-channel", + "mime", + "percent-encoding", + "pin-project-lite", + "rand 0.9.2", + "sha1", + "smallvec", + "tokio", + "tokio-util", + "tracing", + "zstd", +] + +[[package]] +name = "actix-macros" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e01ed3140b2f8d422c68afa1ed2e85d996ea619c988ac834d255db32138655cb" +dependencies = [ + "quote", + "syn 2.0.111", +] + +[[package]] +name = "actix-router" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13d324164c51f63867b57e73ba5936ea151b8a41a1d23d1031eeb9f70d0236f8" +dependencies = [ + "bytestring", + "cfg-if", + "http 0.2.12", + "regex", + "regex-lite", + "serde", + "tracing", +] + +[[package]] +name = "actix-rt" +version = "2.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92589714878ca59a7626ea19734f0e07a6a875197eec751bb5d3f99e64998c63" +dependencies = [ + "futures-core", + "tokio", +] + +[[package]] +name = "actix-server" +version = "2.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a65064ea4a457eaf07f2fba30b4c695bf43b721790e9530d26cb6f9019ff7502" +dependencies = [ + "actix-rt", + "actix-service", + "actix-utils", + "futures-core", + "futures-util", + "mio", + "socket2 0.5.10", + "tokio", + "tracing", +] + +[[package]] +name = "actix-service" +version = "2.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e46f36bf0e5af44bdc4bdb36fbbd421aa98c79a9bce724e1edeb3894e10dc7f" +dependencies = [ + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "actix-utils" +version = "3.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88a1dcdff1466e3c2488e1cb5c36a71822750ad43839937f85d2f4d9f8b705d8" +dependencies = [ + "local-waker", + "pin-project-lite", +] + +[[package]] +name = "actix-web" +version = "4.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1654a77ba142e37f049637a3e5685f864514af11fcbc51cb51eb6596afe5b8d6" +dependencies = [ + "actix-codec", + "actix-http", + "actix-macros", + "actix-router", + "actix-rt", + "actix-server", + "actix-service", + "actix-utils", + "actix-web-codegen", + "bytes", + "bytestring", + "cfg-if", + "cookie", + "derive_more", + "encoding_rs", + "foldhash", + "futures-core", + "futures-util", + "impl-more", + "itoa", + "language-tags", + "log", + "mime", + "once_cell", + "pin-project-lite", + "regex", + "regex-lite", + "serde", + "serde_json", + "serde_urlencoded", + "smallvec", + "socket2 0.6.1", + "time", + "tracing", + "url", +] + +[[package]] +name = "actix-web-codegen" +version = "4.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f591380e2e68490b5dfaf1dd1aa0ebe78d84ba7067078512b4ea6e4492d622b8" +dependencies = [ + "actix-router", + "proc-macro2", + "quote", + "syn 2.0.111", +] + [[package]] name = "addr" version = "0.15.6" @@ -94,6 +277,21 @@ dependencies = [ "equator", ] +[[package]] +name = "alloc-no-stdlib" +version = "2.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3" + +[[package]] +name = "alloc-stdlib" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94fb8275041c72129eb51b7d0322c29b8387a0386127718b096429201a5d6ece" +dependencies = [ + "alloc-no-stdlib", +] + [[package]] name = "allocator-api2" version = "0.2.21" @@ -604,6 +802,27 @@ dependencies = [ "syn 2.0.111", ] +[[package]] +name = "brotli" +version = "8.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4bd8b9603c7aa97359dbd97ecf258968c95f3adddd6db2f7e7a5bef101c84560" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", + "brotli-decompressor", +] + +[[package]] +name = "brotli-decompressor" +version = "5.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "874bb8112abecc98cbd6d81ea4fa7e94fb9449648c93cc89aa40c81c24d7de03" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", +] + [[package]] name = "built" version = "0.8.0" @@ -671,6 +890,15 @@ dependencies = [ "serde", ] +[[package]] +name = "bytestring" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "113b4343b5f6617e7ad401ced8de3cc8b012e73a594347c307b90db3e9271289" +dependencies = [ + "bytes", +] + [[package]] name = "camino" version = "1.2.1" @@ -875,6 +1103,26 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c74b8349d32d297c9134b8c88677813a227df8f779daa29bfc29c183fe3dca6" +[[package]] +name = "convert_case" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "633458d4ef8c78b72454de2d54fd6ab2e60f9e02be22f3c6104cdc8a4e0fceb9" +dependencies = [ + "unicode-segmentation", +] + +[[package]] +name = "cookie" +version = "0.16.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e859cd57d0710d9e06c381b550c06e76992472a8c6d527aecd2fc673dcc231fb" +dependencies = [ + "percent-encoding", + "time", + "version_check", +] + [[package]] name = "core-foundation" version = "0.9.4" @@ -1108,6 +1356,29 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "derive_more" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d751e9e49156b02b44f9c1815bcb94b984cdcc4396ecc32521c739452808b134" +dependencies = [ + "derive_more-impl", +] + +[[package]] +name = "derive_more-impl" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "799a97264921d8623a957f6c3b9011f3b5492f557bbb7a5a19b7fa6d06ba8dcb" +dependencies = [ + "convert_case", + "proc-macro2", + "quote", + "rustc_version", + "syn 2.0.111", + "unicode-xid", +] + [[package]] name = "deunicode" version = "1.6.2" @@ -2239,6 +2510,12 @@ version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e7c5cedc30da3a610cac6b4ba17597bdf7152cf974e8aab3afb3d54455e371c8" +[[package]] +name = "impl-more" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8a5a9a0ff0086c7a148acb942baaabeadf9504d10400b5a05645853729b9cd2" + [[package]] name = "indexmap" version = "1.9.3" @@ -2406,6 +2683,12 @@ dependencies = [ "regex-automata", ] +[[package]] +name = "language-tags" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4345964bb142484797b161f473a503a434de77149dd8c7427788c6e13379388" + [[package]] name = "lazy_static" version = "1.5.0" @@ -2489,6 +2772,23 @@ version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6373607a59f0be73a39b6fe456b8192fcc3585f602af20751600e974dd455e77" +[[package]] +name = "local-channel" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6cbc85e69b8df4b8bb8b89ec634e7189099cea8927a276b7384ce5488e53ec8" +dependencies = [ + "futures-core", + "futures-sink", + "local-waker", +] + +[[package]] +name = "local-waker" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d873d7c67ce09b42110d801813efbc9364414e356be9935700d368351657487" + [[package]] name = "lock_api" version = "0.4.14" @@ -2669,6 +2969,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a69bcab0ad47271a0234d9422b131806bf3968021e5dc9328caf2d4cd58557fc" dependencies = [ "libc", + "log", "wasi", "windows-sys 0.61.2", ] @@ -3713,6 +4014,12 @@ dependencies = [ "regex-syntax", ] +[[package]] +name = "regex-lite" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d942b98df5e658f56f20d592c7f868833fe38115e65c33003d8cd224b0155da" + [[package]] name = "regex-syntax" version = "0.8.8" @@ -5535,6 +5842,12 @@ dependencies = [ "unicode-script", ] +[[package]] +name = "unicode-segmentation" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" + [[package]] name = "unicode-width" version = "0.1.14" @@ -5641,10 +5954,12 @@ name = "void-sentinel" version = "0.1.0-alpha" dependencies = [ "ab_glyph", + "actix-web", "anyhow", "base64 0.22.1", "chrono", "dotenvy", + "futures", "image", "imageproc", "poise", @@ -6403,6 +6718,34 @@ dependencies = [ "syn 2.0.111", ] +[[package]] +name = "zstd" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "7.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f49c4d5f0abb602a93fb8736af2a4f4dd9512e36f7f570d66e65ff867ed3b9d" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.0.16+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748" +dependencies = [ + "cc", + "pkg-config", +] + [[package]] name = "zune-core" version = "0.4.12" diff --git a/Cargo.toml b/Cargo.toml index a9aef6c..f11d041 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,12 +9,13 @@ publish = false [dependencies] ab_glyph = "0.2.32" +actix-web = "4.9" anyhow = "1.0.100" base64 = "0.22.1" chrono = "0.4.42" dotenvy = "0.15.7" futures = "0.3" -image = { version = "0.25.9", features = ["webp"] } +image = { version = "0.25.9", features = ["webp", "jpeg"] } imageproc = "0.25.0" poise = "0.6.1" rand = "0.9.2" diff --git a/example.env b/example.env index 5de0906..b71b0d8 100644 --- a/example.env +++ b/example.env @@ -5,4 +5,9 @@ SURREAL_USER="" SURREAL_PASS="" SURREAL_NS="" SURREAL_DB="" -NVIDIA_API="" \ No newline at end of file +NVIDIA_API="" +OLLAMA_SERVER_URL="" +OLLAMA_MODEL="" +AI_CHAT_COOLDOWN_MS="" +BOT_OWNER_ID="" +IGNORE_RUDE="" \ No newline at end of file diff --git a/src/api.rs b/src/api.rs new file mode 100644 index 0000000..7dbcb30 --- /dev/null +++ b/src/api.rs @@ -0,0 +1,355 @@ +use actix_web::{web, App, HttpServer, HttpResponse, post, get, error, body::EitherBody, dev::{forward_ready, Service, ServiceRequest, ServiceResponse, Transform}}; +use futures::future::LocalBoxFuture; +use futures::FutureExt; +use serde::{Deserialize, Serialize}; +use serenity::cache::Cache; +use std::sync::Arc; +use surrealdb::Surreal; +use surrealdb::engine::remote::ws::Client; +use tracing::{info, warn}; + +#[derive(Deserialize)] +pub struct IsBotThereRequest { + pub guild_ids: Vec, +} + +#[derive(Serialize)] +pub struct IsBotThereResponse { + pub results: Vec, +} + +pub struct ApiState { + pub cache: Arc, + pub db: Surreal, + pub api_key: String, +} + +#[derive(Serialize)] +pub struct LeaderboardMember { + pub user_id: String, + pub username: String, + pub avatar: String, + pub level: u64, + pub xp: u64, + pub rank: usize, +} + +#[derive(Serialize)] +pub struct Role { + pub role_id: String, + pub role_name: String, +} + +#[derive(Serialize)] +pub struct LevelRole { + pub id: String, + pub name: String, +} + +#[derive(Deserialize)] +struct GuildRecord { + level_role_stack: Option>>, +} + +#[derive(Deserialize, Debug)] +struct LeaderboardEntry { + id: surrealdb::sql::Thing, + xp: u64, + level: u64, +} + +/// API Key Authentication Middleware +pub struct ApiKeyMiddleware { + api_key: String, +} + +impl ApiKeyMiddleware { + pub fn new(api_key: String) -> Self { + ApiKeyMiddleware { api_key } + } +} + +impl Transform for ApiKeyMiddleware +where + S: Service, Error = error::Error>, + S::Future: 'static, + B: 'static, +{ + type Response = ServiceResponse>; + type Error = error::Error; + type InitError = (); + type Transform = ApiKeyMiddlewareService; + type Future = futures::future::Ready>; + + fn new_transform(&self, service: S) -> Self::Future { + futures::future::ok(ApiKeyMiddlewareService { + service, + api_key: self.api_key.clone(), + }) + } +} + +pub struct ApiKeyMiddlewareService { + service: S, + api_key: String, +} + +impl Service for ApiKeyMiddlewareService +where + S: Service, Error = error::Error>, + S::Future: 'static, + B: 'static, +{ + type Response = ServiceResponse>; + type Error = error::Error; + type Future = LocalBoxFuture<'static, Result>; + + forward_ready!(service); + + fn call(&self, req: ServiceRequest) -> Self::Future { + let api_key = self.api_key.clone(); + let path = req.path().to_string(); + + // Check API key from header + let header_key = req + .headers() + .get("X-API-Key") + .and_then(|h| h.to_str().ok()) + .unwrap_or("") + .to_string(); + + if header_key.is_empty() || header_key != api_key { + warn!("Unauthorized API request to {} - missing or invalid API key", path); + return Box::pin(async move { + Err(error::ErrorUnauthorized("Missing or invalid API key")) + }); + } + + info!("Authorized API request to {}", path); + + Box::pin( + self.service + .call(req) + .then(|res: Result, error::Error>| async move { + Ok(res?.map_into_left_body()) + }), + ) + } +} + +#[post("/api/is_bot_there")] +async fn is_bot_there( + request_body: web::Json, + data: web::Data, +) -> Result { + info!( + "Processing /api/is_bot_there request with {} guild IDs", + request_body.guild_ids.len() + ); + + let results: Vec = request_body + .guild_ids + .iter() + .map(|guild_id| { + let guild_exists = data.cache.guild(*guild_id).is_some(); + guild_exists + }) + .collect(); + + let found_count = results.iter().filter(|&&b| b).count(); + info!( + "Bot found in {}/{} requested guilds", + found_count, + results.len() + ); + + Ok(HttpResponse::Ok().json(IsBotThereResponse { results })) +} + +#[get("/api/{guild_id}/leaderboard")] +async fn get_leaderboard( + guild_id: web::Path, + data: web::Data, +) -> Result { + let guild_id_value = guild_id.into_inner(); + info!("Processing /api/{}/leaderboard request", guild_id_value); + + // Check if bot is in the guild + if data.cache.guild(guild_id_value).is_none() { + warn!("Bot is not in guild {}", guild_id_value); + return Err(error::ErrorNotFound("Bot is not in this guild")); + } + + // Query all users for this guild, ordered by level and xp + let sql = "SELECT * FROM levels WHERE string::starts_with(record::id(id), $prefix) ORDER BY level DESC, xp DESC"; + let prefix = format!("{}:", guild_id_value); + + let mut response = data.db.query(sql).bind(("prefix", prefix)).await + .map_err(|e| { + warn!("Database query error: {}", e); + error::ErrorInternalServerError("Database query failed") + })?; + + let entries: Vec = response.take(0) + .map_err(|e| { + warn!("Failed to parse database response: {}", e); + error::ErrorInternalServerError("Failed to parse database response") + })?; + + if entries.is_empty() { + info!("No leaderboard data found for guild {}", guild_id_value); + return Ok(HttpResponse::Ok().json(Vec::::new())); + } + + info!("Found {} members in leaderboard for guild {}", entries.len(), guild_id_value); + + // Fetch user data for all entries + let mut leaderboard: Vec = Vec::new(); + + for (i, entry) in entries.iter().enumerate() { + // Extract user id from Surreal Thing + let id_value = &entry.id.id; + let clean_id_str = match id_value { + surrealdb::sql::Id::String(s) => s.as_str().to_string(), + _ => entry.id.id.to_string(), + }; + + let parts: Vec<&str> = clean_id_str.split(':').collect(); + let user_id_str = parts.last().unwrap_or(&"0"); + let user_id_u64 = user_id_str.parse::().unwrap_or(0); + + if user_id_u64 == 0 { + continue; + } + + let mut username = String::from("Unknown User"); + let mut avatar_url = String::new(); + + let user_id = serenity::all::UserId::new(user_id_u64); + + // Try to get member info from cache first + if let Some(guild) = data.cache.guild(guild_id_value) { + if let Some(member) = guild.members.get(&user_id) { + username = member.display_name().to_string(); + avatar_url = member.user.face(); + } + } + + // If not in cache, we'll just use Unknown User + if avatar_url.is_empty() { + info!("User {} not in cache for guild {}", user_id_u64, guild_id_value); + } + + leaderboard.push(LeaderboardMember { + user_id: user_id_u64.to_string(), + username, + avatar: avatar_url, + level: entry.level, + xp: entry.xp, + rank: i + 1, + }); + } + + info!("Returning {} members in leaderboard", leaderboard.len()); + Ok(HttpResponse::Ok().json(leaderboard)) +} + +#[get("/api/{guild_id}/roles")] +async fn get_roles( + guild_id: web::Path, + data: web::Data, +) -> Result { + let guild_id_value = guild_id.into_inner(); + info!("Processing /api/{}/roles request", guild_id_value); + + // Check if bot is in the guild + let guild = data.cache.guild(guild_id_value) + .ok_or_else(|| { + warn!("Bot is not in guild {}", guild_id_value); + error::ErrorNotFound("Bot is not in this guild") + })?; + + // Get all roles from the guild + let roles: Vec = guild.roles + .iter() + .map(|(role_id, role)| Role { + role_id: role_id.to_string(), + role_name: role.name.clone(), + }) + .collect(); + + info!("Returning {} roles for guild {}", roles.len(), guild_id_value); + Ok(HttpResponse::Ok().json(roles)) +} + +#[get("/api/{guild_id}/level/track")] +async fn get_level_tracks( + guild_id: web::Path, + data: web::Data, +) -> Result { + let guild_id_value = guild_id.into_inner(); + info!("Processing /api/{}/level/track request", guild_id_value); + + // Check if bot is in the guild + let guild = data.cache.guild(guild_id_value) + .ok_or_else(|| { + warn!("Bot is not in guild {}", guild_id_value); + error::ErrorNotFound("Bot is not in this guild") + })?; + + // Query the guild record from database + let record: Option = data.db.select(("guilds", guild_id_value.to_string())).await + .map_err(|e| { + warn!("Database query error: {}", e); + error::ErrorInternalServerError("Database query failed") + })?; + + let mut result: std::collections::HashMap> = std::collections::HashMap::new(); + + if let Some(record) = record { + if let Some(level_role_stack) = record.level_role_stack { + for (track_name, role_ids) in level_role_stack { + let mut roles: Vec = Vec::new(); + for role_id in role_ids { + let role_id_obj = serenity::all::RoleId::new(role_id); + let role_name = guild.roles.get(&role_id_obj) + .map(|r| r.name.clone()) + .unwrap_or_else(|| "Unknown Role".to_string()); + + roles.push(LevelRole { + id: role_id.to_string(), + name: role_name, + }); + } + result.insert(track_name, roles); + } + } + } + + info!("Returning {} level tracks for guild {}", result.len(), guild_id_value); + Ok(HttpResponse::Ok().json(result)) +} + +pub async fn start_api_server( + cache: Arc, + db: Surreal, + api_key: String, + port: u16, +) -> std::io::Result<()> { + info!("Starting API server on port {}", port); + + let state = web::Data::new(ApiState { cache, db, api_key: api_key.clone() }); + + HttpServer::new(move || { + App::new() + .app_data(state.clone()) + .wrap(ApiKeyMiddleware::new(api_key.clone())) + .service(is_bot_there) + .service(get_leaderboard) + .service(get_roles) + .service(get_level_tracks) + }) + .bind(("0.0.0.0", port))? + .run() + .await +} diff --git a/src/commands/level.rs b/src/commands/level.rs index a7380b3..1150650 100644 --- a/src/commands/level.rs +++ b/src/commands/level.rs @@ -1,18 +1,42 @@ use crate::{Context, Error}; use ab_glyph::{FontRef, PxScale}; -use image::{ImageBuffer, Rgba}; -use imageproc::drawing::{draw_filled_rect_mut, draw_text_mut}; +use image::{ImageBuffer, Rgba, DynamicImage}; +use imageproc::drawing::{draw_filled_rect_mut, draw_filled_circle_mut, draw_text_mut}; use imageproc::rect::Rect; use poise::serenity_prelude as serenity; use serde::Deserialize; use serenity::Mentionable; use std::collections::HashMap; use std::io::Cursor; +use std::sync::{Arc, OnceLock}; +use std::time::Duration; use serenity::prelude::TypeMapKey; use surrealdb::Surreal; use surrealdb::engine::remote::ws::Client; use futures::future::join_all; +use tokio::sync::RwLock; + +/// Global avatar cache - stores decoded images by URL +/// TTL is handled by periodic cleanup or could use moka crate +static AVATAR_CACHE: OnceLock>>>> = OnceLock::new(); + +fn get_avatar_cache() -> &'static Arc>>> { + AVATAR_CACHE.get_or_init(|| Arc::new(RwLock::new(HashMap::new()))) +} + +/// Global HTTP client for avatar fetches - connection pooling +static AVATAR_CLIENT: OnceLock = OnceLock::new(); + +fn get_avatar_client() -> &'static reqwest::Client { + AVATAR_CLIENT.get_or_init(|| { + reqwest::Client::builder() + .timeout(Duration::from_millis(800)) + .pool_max_idle_per_host(10) + .build() + .unwrap_or_default() + }) +} pub struct DbKey; @@ -516,7 +540,9 @@ struct LeaderboardRenderEntry { level: u64, xp: u64, next_level_xp: u64, - avatar: Option, + avatar: Option>, + /// Used to generate a colored placeholder if avatar is missing + user_id: u64, } #[poise::command(slash_command, prefix_command, guild_only)] @@ -542,13 +568,22 @@ pub async fn leaderboard(ctx: Context<'_>) -> Result<(), Error> { } // 1. Fetch all user data and avatars in parallel + let avatar_client = get_avatar_client(); + let avatar_cache = get_avatar_cache(); + let mut tasks = Vec::new(); + // Try to get guild from cache first (much faster than API calls) + let cached_guild = guild_id.to_guild_cached(&ctx.serenity_context().cache).map(|g| g.clone()); + for (i, entry) in entries.iter().enumerate() { - let ctx = ctx.clone(); // Clone for async + let http = ctx.serenity_context().http.clone(); let entry_level = entry.level; let entry_xp = entry.xp; - let guild_id = guild_id; // capture + let guild_id = guild_id; + let avatar_client = avatar_client.clone(); + let avatar_cache = avatar_cache.clone(); + let cached_guild = cached_guild.clone(); // Extract user id robustly from Surreal `Thing` let id_value = &entry.id.id; @@ -567,38 +602,78 @@ pub async fn leaderboard(ctx: Context<'_>) -> Result<(), Error> { if user_id_u64 != 0 { let user_id = serenity::UserId::new(user_id_u64); - // Prefer guild display name (nickname) when available - if let Ok(member) = ctx.http().get_member(guild_id, user_id).await { + + // Try cache first (instant), then fall back to API + let member_from_cache = cached_guild.as_ref() + .and_then(|g| g.members.get(&user_id).cloned()); + + if let Some(member) = member_from_cache { user_name = member.display_name().to_string(); - let face = member.user.face(); - avatar_url = normalize_avatar_url(&face); - } else { - // Fallback to user object - if let Ok(user) = ctx.http().get_user(user_id).await { - user_name = user - .global_name - .as_ref() - .unwrap_or(&user.name) - .to_string(); - let face = user.face(); - avatar_url = normalize_avatar_url(&face); - } + avatar_url = get_small_avatar_url(&member.user); + } else if let Ok(member) = http.get_member(guild_id, user_id).await { + user_name = member.display_name().to_string(); + avatar_url = get_small_avatar_url(&member.user); + } else if let Ok(user) = http.get_user(user_id).await { + user_name = user.global_name.as_ref().unwrap_or(&user.name).to_string(); + avatar_url = get_small_avatar_url(&user); } - } else { - // keep defaults } - // Fetch avatar image if we have a URL - let mut avatar_img = None; - if !avatar_url.is_empty() { - if let Ok(response) = reqwest::get(&avatar_url).await { - if let Ok(bytes) = response.bytes().await { - if let Ok(img) = image::load_from_memory(&bytes) { - avatar_img = Some(img); + // Check avatar cache first + let avatar_img: Option> = if !avatar_url.is_empty() { + // Fast path: check cache + { + let cache = avatar_cache.read().await; + if let Some(cached) = cache.get(&avatar_url) { + Some(cached.clone()) + } else { + None + } + }.or_else(|| None).map(Some).unwrap_or_else(|| { + // Cache miss - need to fetch (will be done below) + None + }) + } else { + None + }; + + // If not in cache, fetch with short timeout + let avatar_img = if avatar_img.is_some() { + avatar_img + } else if !avatar_url.is_empty() { + // Use tokio timeout for precise control + let fetch_result = tokio::time::timeout( + Duration::from_millis(500), + async { + match avatar_client.get(&avatar_url).send().await { + Ok(response) => response.bytes().await.ok(), + Err(_) => None, } - } + } + ).await; + + match fetch_result { + Ok(Some(bytes)) => { + if let Ok(img) = image::load_from_memory(&bytes) { + let arc_img = Arc::new(img); + // Store in cache for next time + { + let mut cache = avatar_cache.write().await; + // Limit cache size to prevent memory issues + if cache.len() < 500 { + cache.insert(avatar_url, arc_img.clone()); + } + } + Some(arc_img) + } else { + None + } + } + _ => None, } - } + } else { + None + }; LeaderboardRenderEntry { username: user_name, @@ -607,6 +682,7 @@ pub async fn leaderboard(ctx: Context<'_>) -> Result<(), Error> { xp: entry_xp, next_level_xp: (entry_level + 1) * 100, avatar: avatar_img, + user_id: user_id_u64, } }); } @@ -622,7 +698,7 @@ pub async fn leaderboard(ctx: Context<'_>) -> Result<(), Error> { ctx.send( poise::CreateReply::default().attachment(serenity::CreateAttachment::bytes( image_data, - "leaderboard.png", + "leaderboard.jpg", )), ) .await?; @@ -630,30 +706,54 @@ pub async fn leaderboard(ctx: Context<'_>) -> Result<(), Error> { Ok(()) } -fn normalize_avatar_url(url: &str) -> String { - if url.is_empty() { return String::new(); } - // Prefer PNG to ensure decoder compatibility; preserve size if present - // Replace extension .webp -> .png, and enforce format=png when query exists - // Simple approach: if it contains ".webp", swap to ".png"; also add "?size=128" if none - let mut out = url.replace(".webp", ".png"); - if !out.contains("format=") && out.contains("cdn.discordapp.com") { - if out.contains('?') { out.push_str("&format=png"); } else { out.push_str("?format=png"); } +/// Get a small avatar URL (48px) - webp is faster to download & decode +fn get_small_avatar_url(user: &serenity::User) -> String { + // Use webp format (smaller file size, faster download) and small size + // The image crate has webp support enabled + match &user.avatar { + Some(hash) => { + let ext = if hash.is_animated() { "gif" } else { "webp" }; + format!( + "https://cdn.discordapp.com/avatars/{}/{}.{}?size=48", + user.id, hash, ext + ) + } + None => { + // Default avatar - use small size + let index = if let Some(discrim) = user.discriminator { + discrim.get() % 5 + } else { + ((user.id.get() >> 22) % 6) as u16 + }; + format!( + "https://cdn.discordapp.com/embed/avatars/{}.png?size=48", + index + ) + } } - if !out.contains("size=") { - if out.contains('?') { out.push_str("&size=128"); } else { out.push_str("?size=128"); } - } - out } +/// Pre-defined colors for avatar placeholders based on user ID +const PLACEHOLDER_COLORS: [Rgba; 8] = [ + Rgba([114, 137, 218, 255]), // Blurple + Rgba([67, 181, 129, 255]), // Green + Rgba([250, 166, 26, 255]), // Yellow + Rgba([240, 71, 71, 255]), // Red + Rgba([255, 115, 250, 255]), // Pink + Rgba([26, 188, 156, 255]), // Teal + Rgba([230, 126, 34, 255]), // Orange + Rgba([155, 89, 182, 255]), // Purple +]; + fn generate_leaderboard_image( entries: Vec, ) -> Result, Error> { // Image dimensions - let width = 800; + let width = 800u32; let height = 100 + (entries.len() as u32 * 80); // Header + rows let mut image = ImageBuffer::from_pixel(width, height, Rgba([40, 44, 52, 255])); // Dark background - // Load font + // Load font once (compiled into binary) let font_bytes = include_bytes!("../assets/Roboto-Regular.ttf"); let font = FontRef::try_from_slice(font_bytes)?; @@ -683,11 +783,33 @@ fn generate_leaderboard_image( &format!("#{}", entry.rank), ); - // Draw Avatar + // Draw Avatar or colored circle placeholder + let avatar_x = 80i64; + let avatar_y = y_offset as i64 + 16; + if let Some(avatar_img) = &entry.avatar { - let avatar_resized = - avatar_img.resize(60, 60, image::imageops::FilterType::Lanczos3); - image::imageops::overlay(&mut image, &avatar_resized, 80, y_offset as i64 + 10); + // Use Nearest filter - fastest possible, good enough for small avatars + let avatar_resized = avatar_img.resize_exact(48, 48, image::imageops::FilterType::Nearest); + image::imageops::overlay(&mut image, &avatar_resized, avatar_x, avatar_y); + } else { + // Draw colored circle placeholder based on user ID + let color_idx = (entry.user_id % 8) as usize; + let color = PLACEHOLDER_COLORS[color_idx]; + let center_x = (avatar_x + 24) as i32; + let center_y = (avatar_y + 24) as i32; + draw_filled_circle_mut(&mut image, (center_x, center_y), 24, color); + + // Draw first letter of username + let first_char = entry.username.chars().next().unwrap_or('?').to_uppercase().to_string(); + draw_text_mut( + &mut image, + white, + center_x - 8, + center_y - 12, + scale_text, + &font, + &first_char, + ); } // Draw Username @@ -713,14 +835,14 @@ fn generate_leaderboard_image( ); // Draw XP Bar - let bar_width = 300; - let bar_height = 20; - let bar_x = 450; - let bar_y = y_offset + 30; + let bar_width = 300u32; + let bar_height = 20u32; + let bar_x = 450i32; + let bar_y = y_offset as i32 + 30; draw_filled_rect_mut( &mut image, - Rect::at(bar_x as i32, bar_y as i32).of_size(bar_width, bar_height), + Rect::at(bar_x, bar_y).of_size(bar_width, bar_height), bar_bg, ); @@ -734,7 +856,7 @@ fn generate_leaderboard_image( if fill_width > 0 { draw_filled_rect_mut( &mut image, - Rect::at(bar_x as i32, bar_y as i32).of_size(fill_width, bar_height), + Rect::at(bar_x, bar_y).of_size(fill_width, bar_height), bar_fill, ); } @@ -744,16 +866,20 @@ fn generate_leaderboard_image( draw_text_mut( &mut image, white, - bar_x as i32 + 5, - bar_y as i32 + 2, // Centering vertically roughly + bar_x + 5, + bar_y + 2, PxScale::from(14.0), &font, &xp_text, ); } - let mut bytes: Vec = Vec::new(); - image.write_to(&mut Cursor::new(&mut bytes), image::ImageFormat::Png)?; + // Use JPEG for much faster encoding (PNG is slow) + // Quality 85 is a good balance of size vs quality + let mut bytes: Vec = Vec::with_capacity(width as usize * height as usize); + let rgb_image = DynamicImage::ImageRgba8(image).into_rgb8(); + let mut encoder = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut bytes, 85); + encoder.encode_image(&rgb_image)?; Ok(bytes) } diff --git a/src/listener.rs b/src/listener.rs index bcc933a..6c66416 100644 --- a/src/listener.rs +++ b/src/listener.rs @@ -35,6 +35,9 @@ impl EventHandler for Handler { if let Err(e) = crate::commands::utility::process_auto_response(&ctx, &msg).await { tracing::error!("Error processing message for auto-response: {}", e); } + if let Err(e) = crate::commands::fun::handle_ai_chat(&ctx, &msg).await { + tracing::error!("Error processing message for AI chat: {}", e); + } } async fn guild_member_update( diff --git a/src/main.rs b/src/main.rs index 00944b7..7d5296f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,10 +1,11 @@ +mod api; mod commands; mod listener; -use ::serenity::all::GatewayIntents; +use ::serenity::all::{GatewayIntents, UserId}; use dotenvy::dotenv; use poise::{Framework, FrameworkOptions, serenity_prelude as serenity}; -use std::{env, sync::Arc}; +use std::{collections::HashSet, env, sync::Arc}; use surrealdb::Surreal; use surrealdb::engine::remote::ws::{Client, Wss}; use surrealdb::opt::auth::Root; @@ -16,6 +17,7 @@ type Context<'a> = poise::Context<'a, Data, Error>; struct Data { db: Surreal, + ai_chat: Arc, } #[tokio::main()] @@ -32,6 +34,8 @@ async fn main() -> Result<(), Error> { let token = env::var("DISCORD_TOKEN")?; + let api_key = env::var("API_KEY").expect("Expected API_KEY in environment"); + let surreal_address = env::var("SURREAL_ADDRESS").expect("Expected SURREAL_ADDRESS in environment"); let surreal_user = env::var("SURREAL_USER").expect("Expected SURREAL_USER in environment"); @@ -39,6 +43,28 @@ async fn main() -> Result<(), Error> { let surreal_ns = env::var("SURREAL_NS").expect("Expected SURREAL_NS in environment"); let surreal_db = env::var("SURREAL_DB").expect("Expected SURREAL_DB in environment"); + let ollama_url = env::var("OLLAMA_SERVER_URL").expect("Expected OLLAMA_SERVER_URL in environment"); + let ollama_model = env::var("OLLAMA_MODEL").unwrap_or_else(|_| "llama3".to_string()); + let ai_chat_cooldown_ms = env::var("AI_CHAT_COOLDOWN_MS") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(1500); + + let ignore_rude: Vec = env::var("IGNORE_RUDE") + .unwrap_or_default() + .split(',') + .map(|s| s.trim()) + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()) + .collect(); + + let ai_chat_manager = Arc::new(commands::fun::AiChatManager::new( + ollama_url, + ollama_model, + std::time::Duration::from_millis(ai_chat_cooldown_ms), + ignore_rude, + )); + let db = Surreal::new::(&surreal_address).await?; db.signin(Root { @@ -51,8 +77,17 @@ async fn main() -> Result<(), Error> { db.use_ns(&surreal_ns).use_db(&surreal_db).await?; let db_clone = db.clone(); + let ai_chat_clone = ai_chat_manager.clone(); + + let owner_id = env::var("BOT_OWNER_ID") + .expect("Expected BOT_OWNER_ID in environment") + .parse::()?; + let mut owners = HashSet::new(); + owners.insert(UserId::new(owner_id)); + let framework = Framework::builder() .options(FrameworkOptions:: { + owners, commands: vec![ commands::level::set_level_roles(), commands::level::get_level_roles(), @@ -62,10 +97,12 @@ async fn main() -> Result<(), Error> { commands::level::levelup_role_bridger(), commands::fun::say(), commands::fun::urban(), + commands::fun::ai_chat(), commands::utility::auto_response(), commands::utility::view_auto_responses(), commands::utility::delete_auto_response(), commands::utility::edit_auto_response(), + commands::utility::summary(), ], prefix_options: poise::PrefixFrameworkOptions { prefix: Some("!".into()), @@ -80,7 +117,10 @@ async fn main() -> Result<(), Error> { .setup(move |context, _ready, framework| { Box::pin(async move { poise::builtins::register_globally(context, &framework.options().commands).await?; - Ok(Data { db: db_clone }) + Ok(Data { + db: db_clone, + ai_chat: ai_chat_clone, + }) }) }) .build(); @@ -94,10 +134,37 @@ async fn main() -> Result<(), Error> { let mut data = client.data.write().await; data.insert::(db.clone()); data.insert::(db.clone()); + data.insert::(ai_chat_manager.clone()); } - if let Err(why) = client.start_autosharded().await { - eprintln!("An error occurred while running the client: {why}"); + // Get API port from environment, default to 8080 + let api_port = env::var("API_PORT") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(8080); + + // Run bot and API server concurrently + let cache = client.cache.clone(); + let db_for_api = db.clone(); + let bot_task = tokio::spawn(async move { + if let Err(why) = client.start_autosharded().await { + eprintln!("An error occurred while running the client: {why}"); + } + }); + + let api_task = api::start_api_server(cache, db_for_api, api_key, api_port); + + // Wait for API to start, then run bot + tokio::select! { + _ = bot_task => { + info!("Bot task finished"); + } + result = api_task => { + match result { + Ok(_) => info!("API server finished"), + Err(e) => eprintln!("API server error: {}", e), + } + } } Ok(())