From 0db52c5c2b1ec7a0c9dfc876263ca7620c08fd5b Mon Sep 17 00:00:00 2001
From: Natsu Kagami <nki@nkagami.me>
Date: Tue, 31 Dec 2024 05:23:54 +0100
Subject: [PATCH] Allow beatmap command to use just numbers as id

---
 youmubot-osu/src/discord/commands.rs | 87 ++++++++++++++++++----------
 1 file changed, 56 insertions(+), 31 deletions(-)

diff --git a/youmubot-osu/src/discord/commands.rs b/youmubot-osu/src/discord/commands.rs
index dde35a0..27b1177 100644
--- a/youmubot-osu/src/discord/commands.rs
+++ b/youmubot-osu/src/discord/commands.rs
@@ -308,37 +308,7 @@ async fn beatmap<U: HasOsuEnv>(
 
     ctx.defer().await?;
 
-    let beatmap = match map {
-        None => {
-            let Some((BeatmapWithMode(b, mode), bmmods)) =
-                load_beatmap(env, ctx.channel_id(), None as Option<&'_ Message>).await
-            else {
-                return Err(Error::msg("no beatmap mentioned in this channel"));
-            };
-            let mods = bmmods.unwrap_or_else(|| Mods::NOMOD.clone());
-            let info = env
-                .oppai
-                .get_beatmap(b.beatmap_id)
-                .await?
-                .get_possible_pp_with(mode, &mods);
-            EmbedType::Beatmap(Box::new(b), info, mods)
-        }
-        Some(map) => {
-            let Some(results) = stream::select(
-                link_parser::parse_new_links(env, &map),
-                stream::select(
-                    link_parser::parse_old_links(env, &map),
-                    link_parser::parse_short_links(env, &map),
-                ),
-            )
-            .next()
-            .await
-            else {
-                return Err(Error::msg("no beatmap detected in the argument"));
-            };
-            results.embed
-        }
-    };
+    let beatmap = parse_map_input(ctx.channel_id(), env, map, mode).await?;
 
     // override into beatmapset if needed
     let beatmap = if beatmapset == Some(true) {
@@ -429,3 +399,58 @@ fn arg_from_username_or_discord(
         (None, None) => None,
     }
 }
+
+async fn parse_map_input(
+    channel_id: serenity::all::ChannelId,
+    env: &OsuEnv,
+    input: Option<String>,
+    mode: Option<Mode>,
+) -> Result<EmbedType> {
+    Ok(match input {
+        None => {
+            let Some((BeatmapWithMode(b, mode), bmmods)) =
+                load_beatmap(env, channel_id, None as Option<&'_ Message>).await
+            else {
+                return Err(Error::msg("no beatmap mentioned in this channel"));
+            };
+            let mods = bmmods.unwrap_or_else(|| Mods::NOMOD.clone());
+            let info = env
+                .oppai
+                .get_beatmap(b.beatmap_id)
+                .await?
+                .get_possible_pp_with(mode, &mods);
+            EmbedType::Beatmap(Box::new(b), info, mods)
+        }
+        Some(map) => {
+            if let Ok(id) = map.parse::<u64>() {
+                let beatmap = match mode {
+                    None => env.beatmaps.get_beatmap_default(id).await,
+                    Some(mode) => env.beatmaps.get_beatmap(id, mode).await,
+                }?;
+                let info = env
+                    .oppai
+                    .get_beatmap(beatmap.beatmap_id)
+                    .await?
+                    .get_possible_pp_with(beatmap.mode, Mods::NOMOD);
+                return Ok(EmbedType::Beatmap(
+                    Box::new(beatmap),
+                    info,
+                    Mods::NOMOD.clone(),
+                ));
+            }
+            let Some(results) = stream::select(
+                link_parser::parse_new_links(env, &map),
+                stream::select(
+                    link_parser::parse_old_links(env, &map),
+                    link_parser::parse_short_links(env, &map),
+                ),
+            )
+            .next()
+            .await
+            else {
+                return Err(Error::msg("no beatmap detected in the argument"));
+            };
+            results.embed
+        }
+    })
+}