use chrono::{TimeZone, Utc}; use codeforces::{Contest, Problem}; use dashmap::DashMap as HashMap; use lazy_static::lazy_static; use regex::{Captures, Regex}; use serenity::builder::CreateMessage; use serenity::{ builder::CreateEmbed, framework::standard::CommandError, model::channel::Message, utils::MessageBuilder, }; use std::collections::HashMap as StdHashMap; use std::time::Instant; use youmubot_prelude::*; type Client = ::Value; lazy_static! { static ref CONTEST_LINK: Regex = Regex::new( r"https?://codeforces\.com/(contest|gym)s?/(?P\d+)(?:/problem/(?P\w+))?" ) .unwrap(); static ref PROBLEMSET_LINK: Regex = Regex::new( r"https?://codeforces\.com/problemset/problem/(?P\d+)/(?P\w+)" ) .unwrap(); } enum ContestOrProblem { Contest(Contest, Option>), Problem(Problem), } /// Caches the contest list. pub struct ContestCache { contests: HashMap>)>, all_list: RwLock<(StdHashMap, Instant)>, http: Client, } impl TypeMapKey for ContestCache { type Value = ContestCache; } impl ContestCache { /// Creates a new, empty cache. pub(crate) async fn new(http: Client) -> Result { let contests_list = Self::fetch_contest_list(http.clone()).await?; Ok(Self { contests: HashMap::new(), all_list: RwLock::new((contests_list, Instant::now())), http, }) } async fn fetch_contest_list(http: Client) -> Result> { log::info!("Fetching contest list, this might take a few seconds to complete..."); let gyms = Contest::list(&http, true).await?; let contests = Contest::list(&http, false).await?; let r: StdHashMap = gyms .into_iter() .chain(contests.into_iter()) .map(|v| (v.id, v)) .collect(); log::info!("Contest list fetched, {} contests indexed", r.len()); Ok(r) } /// Gets a contest from the cache, fetching from upstream if possible. pub async fn get(&self, contest_id: u64) -> Result<(Contest, Option>)> { if let Some(v) = self.contests.get(&contest_id) { if v.1.is_some() { return Ok(v.clone()); } } self.get_and_store_contest(contest_id).await } async fn get_and_store_contest( &self, contest_id: u64, ) -> Result<(Contest, Option>)> { let (c, p) = match Contest::standings(&self.http, contest_id, |f| f.limit(1, 1)).await { Ok((c, p, _)) => (c, Some(p)), Err(codeforces::Error::Codeforces(s)) if s.ends_with("has not started") => { let c = self.get_from_list(contest_id).await?; (c, None) } Err(v) => return Err(Error::from(v)), }; self.contests.insert(contest_id, (c, p)); Ok(self.contests.get(&contest_id).unwrap().clone()) } async fn get_from_list(&self, contest_id: u64) -> Result { let last_updated = self.all_list.read().await.1; if Instant::now() - last_updated > std::time::Duration::from_secs(60 * 60) { // We update at most once an hour. let mut v = self.all_list.write().await; *v = ( Self::fetch_contest_list(self.http.clone()).await?, Instant::now(), ); } self.all_list .read() .await .0 .get(&contest_id) .cloned() .ok_or_else(|| Error::msg("Contest not found")) } } /// Prints info whenever a problem or contest (or more) is sent on a channel. pub struct InfoHook; #[async_trait] impl Hook for InfoHook { async fn call(&mut self, ctx: &Context, m: &Message) -> Result<()> { if m.author.bot { return Ok(()); } let data = ctx.data.read().await; let contest_cache = data.get::().unwrap(); let matches = parse(&m.content[..], contest_cache) .collect::>() .await; if !matches.is_empty() { m.channel_id .send_message( &ctx, CreateMessage::new() .content("Here are the info of the given Codeforces links!") .embed(print_info_message(&matches[..])), ) .await?; } Ok(()) } } fn parse<'a>( content: &'a str, contest_cache: &'a ContestCache, ) -> impl stream::Stream + 'a { let matches = CONTEST_LINK .captures_iter(content) .chain(PROBLEMSET_LINK.captures_iter(content)) .map(|v| parse_capture(contest_cache, v)) .collect::>() .filter_map(|v| future::ready(v.ok())); matches } fn print_info_message(info: &[(ContestOrProblem, &str)]) -> CreateEmbed { let (problems, contests): (Vec<_>, Vec<_>) = info.iter().partition(|(v, _)| match v { ContestOrProblem::Problem(_) => true, ContestOrProblem::Contest(_, _) => false, }); let mut problems = problems .into_iter() .map(|(v, l)| match v { ContestOrProblem::Problem(p) => (p, l), _ => unreachable!(), }) .collect::>(); let contests = contests .into_iter() .map(|(v, l)| match v { ContestOrProblem::Contest(c, p) => (c, p, l), _ => unreachable!(), }) .collect::>(); problems.sort_by(|(a, _), (b, _)| a.rating.unwrap_or(1500).cmp(&b.rating.unwrap_or(1500))); let mut m = MessageBuilder::new(); if !problems.is_empty() { m.push_line("**Problems**").push_line(""); for (problem, link) in problems { m.push(" - [") .push_bold_safe(format!( "[{}{}] {}", problem.contest_id.unwrap_or(0), problem.index, problem.name )) .push(format!("]({})", link)); if let Some(p) = problem.points { m.push(format!(" | **{:.0}** points", p)); } if let Some(p) = problem.rating { m.push(format!(" | rating **{:.0}**", p)); } if !problem.tags.is_empty() { m.push(format!(" | tags: ||`{}`||", problem.tags.join(", "))); } m.push_line(""); } } m.push_line(""); if !contests.is_empty() { m.push_bold_line("Contests").push_line(""); for (contest, problems, link) in contests { let duration = Duration::from_secs(contest.duration_seconds); m.push(" - [") .push_bold_safe(&contest.name) .push(format!("]({})", link)) .push(format!(" | {}", contest.phase)) .push( problems .as_ref() .map(|v| format!(" | **{}** problems", v.len())) .unwrap_or_else(|| "".to_owned()), ) .push( contest .start_time_seconds .as_ref() .and_then(|v| Utc.timestamp_opt(*v as i64, 0).earliest()) .map(|ts| { format!( " | from {} ({})", ts.format(""), ts.format("") ) }) .unwrap_or_else(|| "".to_owned()), ) .push(format!(" | duration **{}**", duration)); if let Some(p) = &contest.prepared_by { m.push(format!( " | prepared by [{}](https://codeforces.com/profile/{})", p, p )); } m.push_line(""); } } CreateEmbed::new().description(m.build()) } #[allow(clippy::needless_lifetimes)] // Doesn't really work async fn parse_capture<'a>( contest_cache: &ContestCache, cap: Captures<'a>, ) -> Result<(ContestOrProblem, &'a str), CommandError> { let contest_id: u64 = cap .name("contest") .ok_or_else(|| CommandError::from("Contest not captured"))? .as_str() .parse()?; let (contest, problems) = contest_cache.get(contest_id).await?; match cap.name("problem") { Some(p) => { for problem in problems.ok_or_else(|| CommandError::from("Contest hasn't started"))? { if problem.index == p.as_str() { return Ok(( ContestOrProblem::Problem(problem), cap.get(0).unwrap().as_str(), )); } } Err("No such problem in contest".into()) } None => Ok(( ContestOrProblem::Contest(contest, problems), cap.get(0).unwrap().as_str(), )), } }