use codeforces::Contest; use serenity::{ framework::standard::{ macros::{command, group}, Args, CommandResult, }, model::{channel::Message, guild::Member}, utils::MessageBuilder, }; use std::{collections::HashMap, sync::Arc, time::Duration}; use youmubot_prelude::*; mod announcer; mod db; mod embed; mod hook; /// Live-commentating a Codeforces round. mod live; /// The TypeMapKey holding the Client. struct CFClient; impl TypeMapKey for CFClient { type Value = Arc>; } use db::{CfSavedUsers, CfUser}; pub use hook::InfoHook; /// Sets up the CF databases. pub async fn setup(path: &std::path::Path, data: &mut TypeMap, announcers: &mut AnnouncerHandler) { CfSavedUsers::insert_into(data, path.join("cf_saved_users.yaml")) .expect("Must be able to set up DB"); let client = Arc::new(ratelimit::Ratelimit::new( codeforces::Client::new(), 4, std::time::Duration::from_secs(1), )); data.insert::(hook::ContestCache::new(client.clone()).await.unwrap()); data.insert::(client); announcers.add("codeforces", announcer::Announcer); } #[group] #[prefix = "cf"] #[description = "Codeforces-related commands"] #[commands(profile, save, ranks, watch, contestranks)] #[default_command(profile)] pub struct Codeforces; #[command] #[aliases("p", "show", "u", "user", "get")] #[description = "Get an user's profile"] #[usage = "[handle or tag = yourself]"] #[example = "natsukagami"] #[max_args(1)] pub async fn profile(ctx: &Context, m: &Message, mut args: Args) -> CommandResult { let data = ctx.data.read().await; let handle = args .single::() .unwrap_or(UsernameArg::mention(m.author.id)); let http = data.get::().unwrap(); let handle = match handle { UsernameArg::Raw(s) => s, UsernameArg::Tagged(u) => { let db = CfSavedUsers::open(&*data); let user = db.borrow()?.get(&u).map(|u| u.handle.clone()); match user { Some(v) => v, None => { m.reply(&ctx, "no saved account found.").await?; return Ok(()); } } } }; let account = codeforces::User::info(&*http.borrow().await?, &[&handle[..]]) .await? .into_iter() .next(); match account { Some(v) => { m.channel_id .send_message(&ctx, |send| { send.content(format!( "{}: Here is the user that you requested", m.author.mention() )) .embed(|e| embed::user_embed(&v, e)) }) .await } None => m.reply(&ctx, "User not found").await, }?; Ok(()) } #[command] #[description = "Link your Codeforces account to the Discord account, to enjoy Youmu's tracking capabilities."] #[usage = "[handle]"] #[num_args(1)] pub async fn save(ctx: &Context, m: &Message, mut args: Args) -> CommandResult { let data = ctx.data.read().await; let handle = args.single::()?; let http = data.get::().unwrap(); let account = codeforces::User::info(&*http.borrow().await?, &[&handle[..]]) .await? .into_iter() .next(); match account { None => { m.reply(&ctx, "cannot find an account with such handle") .await?; } Some(acc) => { // Collect rating changes data. let rating_changes = acc.rating_changes(&*http.borrow().await?).await?; let mut db = CfSavedUsers::open(&*data); m.reply( &ctx, format!("account `{}` has been linked to your account.", &acc.handle), ) .await?; db.borrow_mut()? .insert(m.author.id, CfUser::save(acc, rating_changes)); } } Ok(()) } #[command] #[description = "See the leaderboard of all people in the server."] #[only_in(guilds)] #[num_args(0)] pub async fn ranks(ctx: &Context, m: &Message) -> CommandResult { let data = ctx.data.read().await; let everyone = { let db = CfSavedUsers::open(&*data); let db = db.borrow()?; db.iter() .map(|(k, v)| (k.clone(), v.clone())) .collect::>() }; let guild = m.guild_id.expect("Guild-only command"); let mut ranks = everyone .into_iter() .map(|(id, cf_user)| { guild .member(&ctx, id) .map(|mem| mem.map(|mem| (mem, cf_user))) }) .collect::>() .filter_map(|v| future::ready(v.ok())) .collect::>() .await; ranks.sort_by(|(_, a), (_, b)| b.rating.unwrap_or(-1).cmp(&a.rating.unwrap_or(-1))); if ranks.is_empty() { m.reply(&ctx, "No saved users in this server.").await?; return Ok(()); } let ranks = Arc::new(ranks); const ITEMS_PER_PAGE: usize = 10; let total_pages = (ranks.len() + ITEMS_PER_PAGE - 1) / ITEMS_PER_PAGE; let last_updated = ranks.iter().map(|(_, cfu)| cfu.last_update).min().unwrap(); paginate_reply_fn( move |page, ctx, msg| { let ranks = ranks.clone(); Box::pin(async move { let page = page as usize; let start = ITEMS_PER_PAGE * page; let end = ranks.len().min(start + ITEMS_PER_PAGE); if start >= end { return Ok(false); } let ranks = &ranks[start..end]; let handle_width = ranks.iter().map(|(_, cfu)| cfu.handle.len()).max().unwrap(); let username_width = ranks .iter() .map(|(mem, _)| mem.distinct().len()) .max() .unwrap(); let mut m = MessageBuilder::new(); m.push_line("```"); // Table header m.push_line(format!( "Rank | Rating | {:hw$} | {:uw$}", "Handle", "Username", hw = handle_width, uw = username_width )); m.push_line(format!( "----------------{:->hw$}---{:->uw$}", "", "", hw = handle_width, uw = username_width )); for (id, (mem, cfu)) in ranks.iter().enumerate() { let id = id + start + 1; m.push_line(format!( "{:>4} | {:>6} | {:hw$} | {:uw$}", format!("#{}", id), cfu.rating .map(|v| v.to_string()) .unwrap_or("----".to_owned()), cfu.handle, mem.distinct(), hw = handle_width, uw = username_width )); } m.push_line("```"); m.push(format!( "Page **{}/{}**. Last updated **{}**", page + 1, total_pages, last_updated.to_rfc2822() )); msg.edit(ctx, |f| f.content(m.build())).await?; Ok(true) }) }, ctx, m, std::time::Duration::from_secs(60), ) .await?; Ok(()) } #[command] #[description = "See the server ranks on a certain contest"] #[usage = "[the contest id]"] #[num_args(1)] #[only_in(guilds)] pub async fn contestranks(ctx: &Context, m: &Message, mut args: Args) -> CommandResult { let data = ctx.data.read().await; let contest_id: u64 = args.single()?; let guild = m.guild_id.unwrap(); // Guild-only command let member_cache = data.get::().unwrap(); let members = CfSavedUsers::open(&*data).borrow()?.clone(); let members = members .into_iter() .map(|(user_id, cf_user)| { member_cache .query(&ctx, user_id, guild) .map(|v| v.map(|v| (cf_user.handle, v))) }) .collect::>() .filter_map(|v| future::ready(v)) .collect::>() .await; let http = data.get::().unwrap(); let (contest, problems, ranks) = Contest::standings(&*http.borrow().await?, contest_id, |f| { f.handles(members.iter().map(|(k, _)| k.clone()).collect()) }) .await?; // Table me let ranks = ranks .into_iter() .flat_map(|v| { v.party .members .iter() .filter_map(|m| { members .get(&m.handle) .cloned() .map(|mem| (mem, m.handle.clone(), v.clone())) }) .collect::>() }) .collect::>(); contest_rank_table(ctx, m, contest, problems, ranks).await?; Ok(()) } #[command] #[description = "Watch a contest and announce any change on the members of the server assigned to the contest."] #[usage = "[the contest id]"] #[num_args(1)] #[required_permissions(MANAGE_CHANNELS)] #[only_in(guilds)] pub async fn watch(ctx: &Context, m: &Message, mut args: Args) -> CommandResult { let contest_id: u64 = args.single()?; live::watch_contest(ctx, m.guild_id.unwrap(), m.channel_id, contest_id).await?; Ok(()) } pub(crate) async fn contest_rank_table( ctx: &Context, reply_to: &Message, contest: Contest, problems: Vec, ranks: Vec<(Member, String, codeforces::RanklistRow)>, ) -> Result<()> { const ITEMS_PER_PAGE: usize = 10; let total_pages = (ranks.len() + ITEMS_PER_PAGE - 1) / ITEMS_PER_PAGE; if ranks.is_empty() { reply_to .reply(&ctx, "No one in this server participated in the contest...") .await?; return Ok(()); } let ranks = Arc::new(ranks); paginate_reply_fn( move |page, ctx, msg| { let contest = contest.clone(); let problems = problems.clone(); let ranks = ranks.clone(); Box::pin(async move { let page = page as usize; let start = page * ITEMS_PER_PAGE; let end = ranks.len().min(start + ITEMS_PER_PAGE); if start >= end { return Ok(false); } let ranks = &ranks[start..end]; let hw = ranks .iter() .map(|(mem, handle, _)| format!("{} ({})", handle, mem.distinct()).len()) .max() .unwrap_or(0) .max(6); let hackw = ranks .iter() .map(|(_, _, row)| { format!( "{}/{}", row.successful_hack_count, row.unsuccessful_hack_count ) .len() }) .max() .unwrap_or(0) .max(5); let mut table = MessageBuilder::new(); let mut header = MessageBuilder::new(); // Header header.push(format!( " Rank | {:hw$} | Total | {:hackw$}", "Handle", "Hacks", hw = hw, hackw = hackw )); for p in &problems { header.push(format!(" | {:4}", p.index)); } let header = header.build(); table .push_line(&header) .push_line(format!("{:-5} | {:5.0} | {: 0.0 { table.push(format!("{:^4.0}", p.points)); } else if let Some(_) = p.best_submission_time_seconds { table.push(format!("{:^4}", "?")); } else if p.rejected_attempt_count > 0 { table.push(format!("{:^4}", format!("-{}", p.rejected_attempt_count))); } else { table.push(format!("{:^4}", "")); } } table.push_line(""); } let mut m = MessageBuilder::new(); m.push_bold_safe(&contest.name) .push(" ") .push_line(contest.url()) .push_codeblock(table.build(), None) .push_line(format!("Page **{}/{}**", page + 1, total_pages)); msg.edit(ctx, |e| e.content(m.build())).await?; Ok(true) }) }, ctx, reply_to, Duration::from_secs(60), ) .await }