Files
void-sentinel/src/commands/level.rs

801 lines
27 KiB
Rust

use crate::{Context, Error};
use ab_glyph::{FontRef, PxScale};
use image::{ImageBuffer, Rgba};
use imageproc::drawing::{draw_filled_rect_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 serenity::prelude::TypeMapKey;
use surrealdb::Surreal;
use surrealdb::engine::remote::ws::Client;
pub struct DbKey;
impl TypeMapKey for DbKey {
type Value = Surreal<Client>;
}
#[derive(Deserialize, serde::Serialize, Debug)]
pub struct UserLevel {
pub xp: u64,
pub level: u64,
pub track: String,
pub last_message: chrono::DateTime<chrono::Utc>,
}
#[derive(Deserialize)]
struct GuildRecord {
level_role_stack: Option<HashMap<String, Vec<u64>>>,
levelup_channel: Option<u64>,
levelup_message: Option<String>,
level_up_role_mapper: Option<HashMap<String, u64>>,
}
#[derive(Deserialize, Debug)]
struct LeaderboardEntry {
id: surrealdb::sql::Thing,
xp: u64,
level: u64,
}
#[poise::command(
slash_command,
prefix_command,
guild_only,
default_member_permissions = "MANAGE_GUILD",
required_permissions = "MANAGE_GUILD"
)]
pub async fn set_level_roles(
ctx: Context<'_>,
#[description = "Name of the level role stack"] name: String,
#[description = "Comma-separated list of role mentions (e.g. @Role1, @Role2, ...)"]
roles: String,
) -> Result<(), Error> {
let guild_id = ctx
.guild_id()
.ok_or_else(|| Error::msg("Guild only command"))?;
let mut role_ids = Vec::new();
for part in roles.split(',') {
let part = part.trim();
if let Some(id_str) = part.strip_prefix("<@&").and_then(|s| s.strip_suffix(">")) {
if let Ok(id) = id_str.parse::<u64>() {
role_ids.push(serenity::RoleId::new(id));
}
} else if let Ok(id) = part.parse::<u64>() {
role_ids.push(serenity::RoleId::new(id));
}
}
if role_ids.is_empty() {
ctx.say("No valid roles found. Please use role mentions (e.g. @Role, @Role2, ...).")
.await?;
return Ok(());
}
let role_count = role_ids.len();
let db = &ctx.data().db;
let role_ids_u64: Vec<u64> = role_ids.iter().map(|r| r.get()).collect();
tracing::info!("Updating level_role_stack for guild {}", guild_id);
let updated: Option<serde::de::IgnoredAny> = db
.update(("guilds", guild_id.to_string()))
.merge(serde_json::json!({
"level_role_stack": {
name.clone(): role_ids_u64
}
}))
.await?;
if updated.is_none() {
tracing::info!("Guild record not found, creating new record");
// Record didn't exist, create it
let _: Option<serde::de::IgnoredAny> = db
.create(("guilds", guild_id.to_string()))
.content(serde_json::json!({
"level_role_stack": {
name.clone(): role_ids_u64
}
}))
.await?;
}
ctx.say(format!(
"Level role stack `{}` updated successfully with {} roles.",
name, role_count
))
.await?;
Ok(())
}
#[poise::command(slash_command, prefix_command, guild_only)]
pub async fn get_level_roles(ctx: Context<'_>) -> Result<(), Error> {
let guild_id = ctx
.guild_id()
.ok_or_else(|| Error::msg("Guild only command"))?;
let db = &ctx.data().db;
let record: Option<GuildRecord> = db.select(("guilds", guild_id.to_string())).await?;
let mut response = String::new();
if let Some(record) = record {
if let Some(stack) = record.level_role_stack {
if stack.is_empty() {
response = "No level role stacks found.".to_string();
} else {
for (name, roles) in stack {
response.push_str(&format!("**{}:** ", name));
let role_mentions: Vec<String> =
roles.iter().map(|id| format!("<@&{}>", id)).collect();
response.push_str(&role_mentions.join(", "));
response.push('\n');
}
}
} else {
response = "No level role stacks found.".to_string();
}
} else {
response = "No configuration found for this guild.".to_string();
}
ctx.send(
poise::CreateReply::default()
.content(response)
.allowed_mentions(serenity::CreateAllowedMentions::new().empty_roles()),
)
.await?;
Ok(())
}
#[poise::command(
slash_command,
prefix_command,
guild_only,
default_member_permissions = "MANAGE_GUILD",
required_permissions = "MANAGE_GUILD"
)]
pub async fn set_levelup_message_channel(
ctx: Context<'_>,
#[description = "Channel to send level up messages to (default: current channel)"]
channel: Option<serenity::Channel>,
) -> Result<(), Error> {
let guild_id = ctx
.guild_id()
.ok_or_else(|| Error::msg("Guild only command"))?;
let db = &ctx.data().db;
let channel_id = match channel {
Some(c) => c.id().get(),
None => ctx.channel_id().get(),
};
let _: Option<serde::de::IgnoredAny> = db
.update(("guilds", guild_id.to_string()))
.merge(serde_json::json!({
"levelup_channel": channel_id
}))
.await?;
ctx.say(format!(
"Level up messages will now be sent to <#{}>.",
channel_id
))
.await?;
Ok(())
}
#[poise::command(
slash_command,
prefix_command,
guild_only,
default_member_permissions = "MANAGE_GUILD",
required_permissions = "MANAGE_GUILD"
)]
pub async fn set_levelup_message(
ctx: Context<'_>,
#[description = "Template message (use {user.mention} and {user.level})"] message: String,
) -> Result<(), Error> {
let guild_id = ctx
.guild_id()
.ok_or_else(|| Error::msg("Guild only command"))?;
let db = &ctx.data().db;
let _: Option<serde::de::IgnoredAny> = db
.update(("guilds", guild_id.to_string()))
.merge(serde_json::json!({
"levelup_message": message
}))
.await?;
ctx.say(format!("Level up message template updated to: {}", message))
.await?;
Ok(())
}
#[poise::command(
slash_command,
prefix_command,
guild_only,
default_member_permissions = "MANAGE_GUILD",
required_permissions = "MANAGE_GUILD"
)]
pub async fn levelup_role_bridger(
ctx: Context<'_>,
#[description = "The role to bridge from"] in_role: serenity::Role,
#[description = "The role to bridge to"] out_role: serenity::Role,
) -> Result<(), Error> {
let guild_id = ctx
.guild_id()
.ok_or_else(|| Error::msg("Guild only command"))?;
let db = &ctx.data().db;
let _: Option<serde::de::IgnoredAny> = db
.update(("guilds", guild_id.to_string()))
.merge(serde_json::json!({
"level_up_role_mapper": {
in_role.id.to_string(): out_role.id.get()
}
}))
.await?;
ctx.say(format!(
"Role bridge created: {} -> {}",
in_role.name, out_role.name
))
.await?;
Ok(())
}
pub async fn process_message(
ctx: &serenity::Context,
msg: &serenity::Message,
) -> Result<(), Error> {
if msg.author.bot {
return Ok(());
}
let guild_id = match msg.guild_id {
Some(id) => id,
None => return Ok(()),
};
let data = ctx.data.read().await;
let db = match data.get::<DbKey>() {
Some(db) => db,
None => {
tracing::error!("Database connection not found in context data");
return Ok(());
}
};
let user_id = msg.author.id;
let record_id = ("levels", format!("{}:{}", guild_id, user_id));
let user_level: Option<UserLevel> = db.select(record_id.clone()).await?;
if let Some(mut level_data) = user_level {
tracing::info!(
"Processing message for user {} in guild {}",
user_id,
guild_id
);
// User exists, add XP
// Check cooldown
let now = chrono::Utc::now();
if now
.signed_duration_since(level_data.last_message)
.num_seconds()
< 60
{
return Ok(());
}
use rand::Rng;
let xp_gain = rand::rng().random_range(15..=25);
level_data.xp += xp_gain;
level_data.last_message = now;
tracing::info!(
"Added {} XP to user {}. Total XP: {}",
xp_gain,
user_id,
level_data.xp
);
let next_level_xp = (level_data.level + 1) * 100;
if level_data.xp >= next_level_xp {
level_data.level += 1;
level_data.xp = 0;
tracing::info!(
"User {} leveled up to {} in guild {}",
user_id,
level_data.level,
guild_id
);
// Handle role assignment
let guild_record: Option<GuildRecord> =
db.select(("guilds", guild_id.to_string())).await?;
let mut role_assigned_id: Option<serenity::RoleId> = None;
if let Some(record) = &guild_record {
if let Some(stack) = &record.level_role_stack {
if let Some(roles) = stack.get(&level_data.track) {
// Check all roles in the stack for level requirements
// We need to fetch role names to check for patterns like "(min - max)"
// Try to get from cache first
let guild = guild_id.to_guild_cached(&ctx.cache).map(|g| g.clone());
// Regex to match "(min - max)"
let re = regex::Regex::new(r"\((\d+)\s*-\s*(\d+)\)").unwrap();
for role_id_u64 in roles {
let role_id = serenity::RoleId::new(*role_id_u64);
// Try to find the role name
let role_name: Option<String> =
if let Some(guild) = &guild {
guild.roles.get(&role_id).map(|r| r.name.clone())
} else {
// Fallback to HTTP if not in cache (expensive but safe)
ctx.http.get_guild_roles(guild_id).await.ok().and_then(
|roles| {
roles
.iter()
.find(|r| r.id == role_id)
.map(|r| r.name.clone())
},
)
};
if let Some(name) = role_name {
if let Some(caps) = re.captures(&name) {
if let (Ok(min), Ok(_max)) =
(caps[1].parse::<u64>(), caps[2].parse::<u64>())
{
// Check if we just hit the min level
if level_data.level == min {
tracing::info!(
"Assigning role {} ({}) to user {}",
role_id,
name,
user_id
);
if let Err(e) = ctx
.http
.add_member_role(
guild_id,
user_id,
role_id,
Some("Level up"),
)
.await
{
tracing::error!("Failed to add role: {}", e);
} else {
role_assigned_id = Some(role_id);
}
}
}
}
}
}
}
}
}
// Determine target channel
let target_channel_id = if let Some(record) = &guild_record {
if let Some(channel_id) = record.levelup_channel {
serenity::ChannelId::new(channel_id)
} else {
msg.channel_id
}
} else {
msg.channel_id
};
// Determine message content
let message_content = if let Some(record) = &guild_record {
if let Some(template) = &record.levelup_message {
let mut content = template
.replace("{user.mention}", &msg.author.mention().to_string())
.replace("{user.level}", &level_data.level.to_string());
if let Some(role_id) = role_assigned_id {
content.push_str(&format!("\nYou also received the <@&{}> role!", role_id));
}
content
} else {
// Default message
if let Some(role_id) = role_assigned_id {
format!(
"Congratulations {}, you leveled up to level {} and received the <@&{}> role!",
msg.author.mention(),
level_data.level,
role_id
)
} else {
format!(
"Congratulations {}, you leveled up to level {}!",
msg.author.mention(),
level_data.level
)
}
}
} else {
// Default message if no record
if let Some(role_id) = role_assigned_id {
format!(
"Congratulations {}, you leveled up to level {} and received the <@&{}> role!",
msg.author.mention(),
level_data.level,
role_id
)
} else {
format!(
"Congratulations {}, you leveled up to level {}!",
msg.author.mention(),
level_data.level
)
}
};
let message = serenity::CreateMessage::new()
.content(message_content)
.allowed_mentions(serenity::CreateAllowedMentions::new().users(vec![user_id]));
if let Err(e) = target_channel_id.send_message(&ctx.http, message).await {
tracing::error!("Failed to send level up message: {}", e);
}
}
let _: Option<UserLevel> = db.update(record_id).content(level_data).await?;
} else {
// User doesn't exist, check if they have a starting role
let guild_record: Option<GuildRecord> = db.select(("guilds", guild_id.to_string())).await?;
if let Some(record) = guild_record {
if let Some(stack) = record.level_role_stack {
let member = match guild_id.member(&ctx.http, user_id).await {
Ok(m) => m,
Err(_) => return Ok(()),
};
for (track_name, roles) in stack {
if let Some(first_role_id) = roles.first() {
if member
.roles
.contains(&serenity::RoleId::new(*first_role_id))
{
// User has the starting role for this track
let new_level = UserLevel {
xp: 0,
level: 0,
track: track_name.clone(),
last_message: chrono::Utc::now(),
};
tracing::info!(
"Initializing user {} on track {} in guild {}",
user_id,
track_name,
guild_id
);
let _: Option<UserLevel> =
db.create(record_id).content(new_level).await?;
break; // Only assign to one track
}
}
}
}
}
}
Ok(())
}
#[poise::command(slash_command, prefix_command, guild_only)]
pub async fn leaderboard(ctx: Context<'_>) -> Result<(), Error> {
let guild_id = ctx
.guild_id()
.ok_or_else(|| Error::msg("Guild only command"))?;
let db = &ctx.data().db;
// Defer interaction as image generation might take time
ctx.defer().await?;
// Query top 10 users for this guild
// We filter by ID starting with "levels:guild_id:"
// Note: In SurrealDB, record IDs are `table:id_part`.
// The id_part here is `guild_id:user_id`.
// We can use string functions on the ID.
let sql = "SELECT * FROM levels WHERE string::starts_with(record::id(id), $prefix) ORDER BY level DESC, xp DESC LIMIT 10";
let prefix = format!("{}:", guild_id);
let mut response = db.query(sql).bind(("prefix", prefix)).await?;
let entries: Vec<LeaderboardEntry> = response.take(0)?;
if entries.is_empty() {
ctx.say("No leaderboard data found for this guild.").await?;
return Ok(());
}
// Generate image
let image_data = generate_leaderboard_image(&ctx, &entries).await?;
ctx.send(
poise::CreateReply::default().attachment(serenity::CreateAttachment::bytes(
image_data,
"leaderboard.png",
)),
)
.await?;
Ok(())
}
async fn generate_leaderboard_image(
ctx: &Context<'_>,
entries: &[LeaderboardEntry],
) -> Result<Vec<u8>, Error> {
// Image dimensions
let width = 800;
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
let font_bytes = include_bytes!("../assets/Roboto-Regular.ttf");
let font = FontRef::try_from_slice(font_bytes)?;
let white = Rgba([255, 255, 255, 255]);
let gray = Rgba([180, 180, 180, 255]);
let bar_bg = Rgba([60, 63, 69, 255]);
let bar_fill = Rgba([114, 137, 218, 255]); // Blurple-ish
// Draw Title
let scale_title = PxScale::from(40.0);
draw_text_mut(&mut image, white, 20, 20, scale_title, &font, "Leaderboard");
let scale_text = PxScale::from(24.0);
let scale_small = PxScale::from(18.0);
for (i, entry) in entries.iter().enumerate() {
let y_offset = 100 + (i as u32 * 80);
// Extract user ID from record ID
// id.id is String("guild_id:user_id")
let id_str = entry.id.id.to_string();
let clean_id_str = id_str.trim_matches(|c| c == '"' || c == '⟨' || c == '⟩');
let parts: Vec<&str> = clean_id_str.split(':').collect();
let user_id_str = parts.last().unwrap_or(&"0");
let user_id = user_id_str.parse::<u64>().unwrap_or(0);
// Fetch user info
let user = if user_id != 0 {
match ctx.http().get_user(serenity::UserId::new(user_id)).await {
Ok(u) => u,
Err(_) => {
// Fallback if user not found
let mut u = serenity::User::default();
u.name = "Unknown User".to_string();
u
}
}
} else {
tracing::warn!("Invalid user ID parsed from {}: {}", id_str, user_id_str);
let mut u = serenity::User::default();
u.name = "Unknown User".to_string();
u
};
// Draw Rank
draw_text_mut(
&mut image,
white,
20,
y_offset as i32 + 20,
scale_text,
&font,
&format!("#{}", i + 1),
);
// Draw Avatar
let avatar_url = user.face();
// We need to fetch the avatar image
if let Ok(response) = reqwest::get(&avatar_url).await {
if let Ok(bytes) = response.bytes().await {
if let Ok(avatar_img) = image::load_from_memory(&bytes) {
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);
}
}
}
// Draw Username
draw_text_mut(
&mut image,
white,
160,
y_offset as i32 + 15,
scale_text,
&font,
&user.name,
);
// Draw Level
draw_text_mut(
&mut image,
gray,
160,
y_offset as i32 + 45,
scale_small,
&font,
&format!("Level {}", entry.level),
);
// Draw XP Bar
let bar_width = 300;
let bar_height = 20;
let bar_x = 450;
let bar_y = y_offset + 30;
draw_filled_rect_mut(
&mut image,
Rect::at(bar_x as i32, bar_y as i32).of_size(bar_width, bar_height),
bar_bg,
);
let next_level_xp = (entry.level + 1) * 100;
let progress = if next_level_xp > 0 {
entry.xp as f32 / next_level_xp as f32
} else {
0.0
};
let fill_width = (progress * bar_width as f32) as u32;
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),
bar_fill,
);
}
// Draw XP Text
let xp_text = format!("{}/{} XP", entry.xp, next_level_xp);
draw_text_mut(
&mut image,
white,
bar_x as i32 + 5,
bar_y as i32 + 2, // Centering vertically roughly
PxScale::from(14.0),
&font,
&xp_text,
);
}
let mut bytes: Vec<u8> = Vec::new();
image.write_to(&mut Cursor::new(&mut bytes), image::ImageFormat::Png)?;
Ok(bytes)
}
pub async fn on_guild_member_update(
ctx: &serenity::Context,
_old: &Option<serenity::Member>,
new: &serenity::Member,
) -> Result<(), Error> {
let guild_id = new.guild_id;
let data = ctx.data.read().await;
let db = match data.get::<DbKey>() {
Some(db) => db,
None => {
tracing::error!("Database connection not found in context data");
return Ok(());
}
};
let guild_record: Option<GuildRecord> = db.select(("guilds", guild_id.to_string())).await?;
if let Some(record) = guild_record {
if let Some(mapper) = record.level_up_role_mapper {
// Check if user has any "in_role"
// We iterate through the mapper to find if any key (in_role) matches user's roles
for (in_role_str, out_role_id) in &mapper {
if let Ok(in_role_id) = in_role_str.parse::<u64>() {
let in_role = serenity::RoleId::new(in_role_id);
if new.roles.contains(&in_role) {
// User has the in_role.
// Now check if they have ANY other out_role from the mapper
let mut has_any_out_role = false;
for (_, other_out_role_id) in &mapper {
if new
.roles
.contains(&serenity::RoleId::new(*other_out_role_id))
{
has_any_out_role = true;
break;
}
}
if !has_any_out_role {
// User has no out_roles, so we bridge them
tracing::info!(
"Bridging user {} from role {} to {}",
new.user.id,
in_role_id,
out_role_id
);
// Add out_role
if let Err(e) = ctx
.http
.add_member_role(
guild_id,
new.user.id,
serenity::RoleId::new(*out_role_id),
Some("Role Bridge"),
)
.await
{
tracing::error!("Failed to add bridged role: {}", e);
}
// Remove in_role
if let Err(e) = ctx
.http
.remove_member_role(
guild_id,
new.user.id,
in_role,
Some("Role Bridge"),
)
.await
{
tracing::error!("Failed to remove bridged role: {}", e);
}
} else {
// User already has an out_role, just remove the in_role
tracing::info!(
"User {} already has an out_role, removing in_role {}",
new.user.id,
in_role_id
);
if let Err(e) = ctx
.http
.remove_member_role(
guild_id,
new.user.id,
in_role,
Some("Role Bridge cleanup"),
)
.await
{
tracing::error!("Failed to remove bridged role: {}", e);
}
}
}
}
}
}
}
Ok(())
}