From f3c062f417125c0b6c5acce7e6543868f6330d69 Mon Sep 17 00:00:00 2001
From: Nguyen Duc Huy <ndhuy01@gmail.com>
Date: Sun, 17 Mar 2024 17:13:57 -0400
Subject: [PATCH] Getting pp info when a new user is saved. (#42)

---
 Cargo.lock                      |  1 +
 youmubot-osu/Cargo.toml         |  1 +
 youmubot-osu/src/discord/mod.rs | 41 ++++++++++++++++++++++++++++++---
 3 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index d198d81..100d036 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3181,6 +3181,7 @@ dependencies = [
  "bitflags 1.3.2",
  "chrono",
  "dashmap",
+ "futures-util",
  "lazy_static",
  "osuparse",
  "rand",
diff --git a/youmubot-osu/Cargo.toml b/youmubot-osu/Cargo.toml
index 53b17fe..adcf309 100644
--- a/youmubot-osu/Cargo.toml
+++ b/youmubot-osu/Cargo.toml
@@ -22,6 +22,7 @@ serde = { version = "1.0.137", features = ["derive"] }
 serenity = "0.12"
 zip = "0.6.2"
 rand = "0.8"
+futures-util = "0.3.30"
 
 youmubot-db = { path = "../youmubot-db" }
 youmubot-db-sql = { path = "../youmubot-db-sql" }
diff --git a/youmubot-osu/src/discord/mod.rs b/youmubot-osu/src/discord/mod.rs
index 30eb979..3c1a32b 100644
--- a/youmubot-osu/src/discord/mod.rs
+++ b/youmubot-osu/src/discord/mod.rs
@@ -1,5 +1,6 @@
 use std::{str::FromStr, sync::Arc};
 
+use futures_util::join;
 use rand::seq::IteratorRandom;
 use serenity::{
     builder::{CreateMessage, EditMessage},
@@ -80,7 +81,6 @@ impl TypeMapKey for OsuEnv {
 ///  This does NOT automatically enable:
 ///  - Commands on the "osu" prefix
 ///  - Hooks. Hooks are completely opt-in.
-///  
 pub async fn setup(
     data: &mut TypeMap,
     prelude: youmubot_prelude::Env,
@@ -364,14 +364,49 @@ pub async fn forcesave(ctx: &Context, msg: &Message, mut args: Args) -> CommandR
 }
 
 async fn add_user(target: serenity::model::id::UserId, user: User, env: &OsuEnv) -> Result<()> {
+    let pp_fut = async {
+        [Mode::Std, Mode::Taiko, Mode::Catch, Mode::Mania]
+            .into_iter()
+            .map(|mode| async move {
+                env.client
+                    .user(UserID::ID(user.id), |f| f.mode(mode))
+                    .await
+                    .unwrap_or_else(|err| {
+                        eprintln!("{}", err);
+                        None
+                    })
+                    .and_then(|u| u.pp)
+            })
+            .collect::<stream::FuturesOrdered<_>>()
+            .collect::<Vec<_>>()
+            .await
+    };
+
+    let std_weight_map_length_fut = async {
+        let scores = env
+            .client
+            .user_best(UserID::ID(user.id), |f| f.mode(Mode::Std).limit(100))
+            .await
+            .unwrap_or_else(|err| {
+                eprintln!("{}", err);
+                vec![]
+            });
+
+        calculate_weighted_map_length(&scores, &env.beatmaps, Mode::Std)
+            .await
+            .pls_ok()
+    };
+
+    let (pp, std_weight_map_length) = join!(pp_fut, std_weight_map_length_fut);
+
     let u = OsuUser {
         user_id: target,
         username: user.username.into(),
         id: user.id,
         failures: 0,
         last_update: chrono::Utc::now(),
-        pp: [None, None, None, None],
-        std_weighted_map_length: None,
+        pp: pp.try_into().unwrap(),
+        std_weighted_map_length: std_weight_map_length,
     };
     env.saved_users.new_user(u).await?;
     Ok(())