use crate::{ discord::beatmap_cache::BeatmapMetaCache, discord::oppai_cache::{BeatmapCache, BeatmapInfoWithPP}, models::{Beatmap, Mode, Mods}, }; use lazy_static::lazy_static; use regex::Regex; use serenity::{builder::CreateEmbed, model::channel::Message, utils::MessageBuilder}; use std::str::FromStr; use youmubot_prelude::*; use super::embeds::beatmap_embed; lazy_static! { pub(crate) static ref OLD_LINK_REGEX: Regex = Regex::new( r"(?:https?://)?osu\.ppy\.sh/(?Ps|b)/(?P\d+)(?:[\&\?]m=(?P\d))?(?:\+(?P[A-Z]+))?" ).unwrap(); pub(crate) static ref NEW_LINK_REGEX: Regex = Regex::new( r"(?:https?://)?osu\.ppy\.sh/beatmapsets/(?P\d+)/?(?:\#(?Posu|taiko|fruits|mania)(?:/(?P\d+)|/?))?(?:\+(?P[A-Z]+))?" ).unwrap(); pub(crate) static ref SHORT_LINK_REGEX: Regex = Regex::new( r"(?:^|\s|\W)(?P
/b/(?P\d+)(?:/(?Posu|taiko|fruits|mania))?(?:\+(?P[A-Z]+))?)" ).unwrap(); } pub fn dot_osu_hook<'a>( ctx: &'a Context, msg: &'a Message, ) -> std::pin::Pin> + Send + 'a>> { Box::pin(async move { if msg.author.bot { return Ok(()); } // Take all the .osu attachments let mut osu_embeds = msg .attachments .iter() .filter( |a| a.filename.ends_with(".osu") && a.size < 1024 * 1024, /* 1mb */ ) .map(|attachment| { let url = attachment.url.clone(); async move { let data = ctx.data.read().await; let oppai = data.get::().unwrap(); let (beatmap, _) = oppai.download_beatmap_from_url(&url).await.ok()?; let embed_fn = crate::discord::embeds::beatmap_offline_embed( &beatmap, Mode::from(beatmap.content.mode as u8), /*For now*/ msg.content.trim().parse().unwrap_or(Mods::NOMOD), ) .ok()?; let mut create_embed = CreateEmbed::default(); embed_fn(&mut create_embed); Some(create_embed) } }) .collect::>() .filter_map(future::ready) .collect::>() .await; let osz_embeds = msg .attachments .iter() .filter( |a| a.filename.ends_with(".osz") && a.size < 20 * 1024 * 1024, /* 20mb */ ) .map(|attachment| { let url = attachment.url.clone(); async move { let data = ctx.data.read().await; let oppai = data.get::().unwrap(); let beatmaps = oppai.download_osz_from_url(&url).await.pls_ok()?; Some( beatmaps .into_iter() .filter_map(|beatmap| { let embed_fn = crate::discord::embeds::beatmap_offline_embed( &beatmap, Mode::from(beatmap.content.mode as u8), /*For now*/ msg.content.trim().parse().unwrap_or(Mods::NOMOD), ) .pls_ok()?; let mut create_embed = CreateEmbed::default(); embed_fn(&mut create_embed); Some(create_embed) }) .collect::>(), ) } }) .collect::>() .filter_map(future::ready) .filter(|v| future::ready(!v.is_empty())) .collect::>() .await .concat(); osu_embeds.extend(osz_embeds); if !osu_embeds.is_empty() { msg.channel_id .send_message(ctx, |f| { f.reference_message(msg) .content(format!("{} attached beatmaps found", osu_embeds.len())) .add_embeds(osu_embeds) }) .await .ok(); } Ok(()) }) } pub fn hook<'a>( ctx: &'a Context, msg: &'a Message, ) -> std::pin::Pin> + Send + 'a>> { Box::pin(async move { if msg.author.bot { return Ok(()); } let (old_links, new_links, short_links) = ( handle_old_links(ctx, &msg.content), handle_new_links(ctx, &msg.content), handle_short_links(ctx, msg, &msg.content), ); stream::select(old_links, stream::select(new_links, short_links)) .then(|l| async move { match l.embed { EmbedType::Beatmap(b, info, mods) => { handle_beatmap(ctx, &b, info, l.link, l.mode, mods, msg) .await .pls_ok(); let mode = l.mode.unwrap_or(b.mode); let bm = super::BeatmapWithMode(*b, mode); crate::discord::cache::save_beatmap( &*ctx.data.read().await, msg.channel_id, &bm, ) .await .pls_ok(); } EmbedType::Beatmapset(b) => { handle_beatmapset(ctx, b, l.link, l.mode, msg) .await .pls_ok(); } } }) .collect::<()>() .await; Ok(()) }) } enum EmbedType { Beatmap(Box, BeatmapInfoWithPP, Mods), Beatmapset(Vec), } struct ToPrint<'a> { embed: EmbedType, link: &'a str, mode: Option, } fn handle_old_links<'a>( ctx: &'a Context, content: &'a str, ) -> impl stream::Stream> + 'a { OLD_LINK_REGEX .captures_iter(content) .map(move |capture| async move { let data = ctx.data.read().await; let cache = data.get::().unwrap(); let osu = data.get::().unwrap(); let req_type = capture.name("link_type").unwrap().as_str(); let mode = capture .name("mode") .map(|v| v.as_str().parse()) .transpose()? .and_then(|v| { Some(match v { 0 => Mode::Std, 1 => Mode::Taiko, 2 => Mode::Catch, 3 => Mode::Mania, _ => return None, }) }); let beatmaps = match req_type { "b" => vec![match mode { Some(mode) => osu.get_beatmap(capture["id"].parse()?, mode).await?, None => osu.get_beatmap_default(capture["id"].parse()?).await?, }], "s" => osu.get_beatmapset(capture["id"].parse()?).await?, _ => unreachable!(), }; if beatmaps.is_empty() { return Ok(None); } let r: Result<_> = Ok(match req_type { "b" => { let b = Box::new(beatmaps.into_iter().next().unwrap()); // collect beatmap info let mods = capture .name("mods") .and_then(|v| Mods::from_str(v.as_str()).pls_ok()) .unwrap_or(Mods::NOMOD); let info = { let mode = mode.unwrap_or(b.mode); cache .get_beatmap(b.beatmap_id) .await .and_then(|b| b.get_possible_pp_with(mode, mods))? }; Some(ToPrint { embed: EmbedType::Beatmap(b, info, mods), link: capture.get(0).unwrap().as_str(), mode, }) } "s" => Some(ToPrint { embed: EmbedType::Beatmapset(beatmaps), link: capture.get(0).unwrap().as_str(), mode, }), _ => None, }); r }) .collect::>() .filter_map(|v| { future::ready(match v { Ok(v) => v, Err(e) => { eprintln!("{}", e); None } }) }) } fn handle_new_links<'a>( ctx: &'a Context, content: &'a str, ) -> impl stream::Stream> + 'a { NEW_LINK_REGEX .captures_iter(content) .map(|capture| async move { let data = ctx.data.read().await; let osu = data.get::().unwrap(); let cache = data.get::().unwrap(); let mode = capture .name("mode") .and_then(|v| Mode::parse_from_new_site(v.as_str())); let link = capture.get(0).unwrap().as_str(); let beatmaps = match capture.name("beatmap_id") { Some(ref v) => vec![match mode { Some(mode) => osu.get_beatmap(v.as_str().parse()?, mode).await?, None => osu.get_beatmap_default(v.as_str().parse()?).await?, }], None => { osu.get_beatmapset(capture.name("set_id").unwrap().as_str().parse()?) .await? } }; if beatmaps.is_empty() { return Ok(None); } let r: Result<_> = Ok(match capture.name("beatmap_id") { Some(_) => { let beatmap = Box::new(beatmaps.into_iter().next().unwrap()); // collect beatmap info let mods = capture .name("mods") .and_then(|v| Mods::from_str(v.as_str()).pls_ok()) .unwrap_or(Mods::NOMOD); let info = { let mode = mode.unwrap_or(beatmap.mode); cache .get_beatmap(beatmap.beatmap_id) .await .and_then(|b| b.get_possible_pp_with(mode, mods))? }; Some(ToPrint { embed: EmbedType::Beatmap(beatmap, info, mods), link, mode, }) } None => Some(ToPrint { embed: EmbedType::Beatmapset(beatmaps), link, mode, }), }); r }) .collect::>() .filter_map(|v| { future::ready(match v { Ok(v) => v, Err(e) => { eprintln!("{}", e); None } }) }) } fn handle_short_links<'a>( ctx: &'a Context, msg: &'a Message, content: &'a str, ) -> impl stream::Stream> + 'a { SHORT_LINK_REGEX .captures_iter(content) .map(|capture| async move { if let Some(guild_id) = msg.guild_id { if announcer::announcer_of(ctx, crate::discord::announcer::ANNOUNCER_KEY, guild_id) .await? != Some(msg.channel_id) { // Disable if we are not in the server's announcer channel return Err(Error::msg("not in server announcer channel")); } } let data = ctx.data.read().await; let osu = data.get::().unwrap(); let cache = data.get::().unwrap(); let mode = capture .name("mode") .and_then(|v| Mode::parse_from_new_site(v.as_str())); let id: u64 = capture.name("id").unwrap().as_str().parse()?; let beatmap = match mode { Some(mode) => osu.get_beatmap(id, mode).await, None => osu.get_beatmap_default(id).await, }?; let mods = capture .name("mods") .and_then(|v| Mods::from_str(v.as_str()).pls_ok()) .unwrap_or(Mods::NOMOD); let info = { let mode = mode.unwrap_or(beatmap.mode); cache .get_beatmap(beatmap.beatmap_id) .await .and_then(|b| b.get_possible_pp_with(mode, mods))? }; let r: Result<_> = Ok(ToPrint { embed: EmbedType::Beatmap(Box::new(beatmap), info, mods), link: capture.name("main").unwrap().as_str(), mode, }); r }) .collect::>() .filter_map(|v| { future::ready(match v { Ok(v) => Some(v), Err(e) => { eprintln!("{}", e); None } }) }) } async fn handle_beatmap<'a, 'b>( ctx: &Context, beatmap: &Beatmap, info: BeatmapInfoWithPP, link: &'_ str, mode: Option, mods: Mods, reply_to: &Message, ) -> Result<()> { reply_to .channel_id .send_message(ctx, |m| { m.content( MessageBuilder::new() .push("Beatmap information for ") .push_mono_safe(link) .build(), ) .embed(|b| beatmap_embed(beatmap, mode.unwrap_or(beatmap.mode), mods, info, b)) .reference_message(reply_to) }) .await?; Ok(()) } async fn handle_beatmapset<'a, 'b>( ctx: &Context, beatmaps: Vec, link: &'_ str, mode: Option, reply_to: &Message, ) -> Result<()> { crate::discord::display::display_beatmapset( ctx, beatmaps, mode, None, reply_to, format!("Beatmapset information for `{}`", link), ) .await .pls_ok(); Ok(()) }