feat: cache bttv emotes on joinHEADmaster

1 files changed, 42 insertions, 9 deletions

diff --git a/src/betterttv.rs b/src/betterttv.rs
index 09f544c..8f8f3e7 100644
--- a/src/betterttv.rs
+++ b/src/betterttv.rs
@@ -1,4 +1,7 @@
-use std::{collections::HashSet, sync::Arc};
+use std::{
+    collections::{HashMap, HashSet},
+    sync::Arc,
+};
 
 use futures::SinkExt;
 use reqwest::{Client, Error};
@@ -95,6 +98,8 @@ impl RetrieveEmoteAPI<BetterTTVEmote> for BetterTTVAPIClient {
 }
 
 pub struct BetterTTVWSClient {
+    api_client: BetterTTVAPIClient,
+
     url: String,
     on_emote_create: Option<RetrieveEmoteHandler<Emote>>,
     on_emote_update: Option<RetrieveEmoteHandler<Emote>>,
@@ -102,6 +107,8 @@ pub struct BetterTTVWSClient {
 
     joined_channels: HashSet<usize>,
     awaiting_channels: HashSet<usize>,
+
+    emotes: HashMap<String, HashMap<String, String>>,
 }
 
 impl RetrieveEmoteWS<Emote> for BetterTTVWSClient {
@@ -136,12 +143,14 @@ impl BetterTTVWSClient {
         Ok((
             s.0,
             Self {
+                api_client: BetterTTVAPIClient::new(),
                 url: url.to_string(),
                 on_emote_create: None,
                 on_emote_delete: None,
                 on_emote_update: None,
                 joined_channels: HashSet::new(),
                 awaiting_channels: HashSet::new(),
+                emotes: HashMap::new(),
             },
         ))
     }
@@ -184,7 +193,7 @@ impl BetterTTVWSClient {
         self.awaiting_channels.insert(twitch_id);
     }
 
-    fn process_message(&self, msg: Message) {
+    fn process_message(&mut self, msg: Message) {
         match msg {
             Message::Text(text) => {
                 let text = text.to_string();
@@ -195,15 +204,20 @@ impl BetterTTVWSClient {
                 let event_data = &json["data"];
 
                 let event_name = json["name"].as_str().expect("No event name");
+                let channel_name = event_data["channel"]
+                    .as_str()
+                    .expect("No channel")
+                    .to_string()
+                    .clone();
+
+                let Some(emote_cache) = self.emotes.get_mut(&channel_name) else {
+                    return;
+                };
 
                 if event_name.eq("emote_create") {
                     if let Some(func) = &self.on_emote_create {
                         let emote_data = &event_data["emote"];
-                        let channel_data = event_data["channel"]
-                            .as_str()
-                            .expect("No channel")
-                            .to_string()
-                            .clone();
+                        let channel_data = channel_name;
 
                         let emote = Emote {
                             id: emote_data["id"].as_str().expect("No emote.id").to_string(),
@@ -214,6 +228,8 @@ impl BetterTTVWSClient {
                             original_code: None,
                         };
 
+                        emote_cache.insert(emote.id.clone(), emote.code.clone());
+
                         (func)(channel_data, None, emote);
                     }
                 } else if event_name.eq("emote_update") {
@@ -225,7 +241,7 @@ impl BetterTTVWSClient {
                             .to_string()
                             .clone();
 
-                        let emote = Emote {
+                        let mut emote = Emote {
                             id: emote_data["id"].as_str().expect("No emote.id").to_string(),
                             code: emote_data["code"]
                                 .as_str()
@@ -234,6 +250,9 @@ impl BetterTTVWSClient {
                             original_code: None,
                         };
 
+                        emote.original_code = emote_cache.get(&emote.id).cloned();
+                        emote_cache.insert(emote.id.clone(), emote.code.clone());
+
                         (func)(channel_data, None, emote);
                     }
                 } else if event_name.eq("emote_delete") {
@@ -245,12 +264,15 @@ impl BetterTTVWSClient {
                             .to_string()
                             .clone();
 
-                        let emote = Emote {
+                        let mut emote = Emote {
                             id: emote_id.as_str().expect("No emoteId").to_string(),
                             code: "".into(),
                             original_code: None,
                         };
 
+                        emote.code = emote_cache.get(&emote.id).cloned().unwrap_or("".into());
+                        emote_cache.remove(&emote.id);
+
                         (func)(channel_data, None, emote);
                     }
                 }
@@ -263,6 +285,7 @@ impl BetterTTVWSClient {
         for id in &self.awaiting_channels {
             let json = if self.joined_channels.contains(id) {
                 self.joined_channels.remove(id);
+                self.emotes.remove(&format!("twitch:{}", id));
 
                 serde_json::json!({
                     "name": "part_channel",
@@ -273,6 +296,16 @@ impl BetterTTVWSClient {
             } else {
                 self.joined_channels.insert(*id);
 
+                // Caching channel emotes
+                let mut emote_map = HashMap::new();
+                if let Ok(emotes) = self.api_client.get_channel_emotes(&id.to_string()).await {
+                    for emote in emotes {
+                        emote_map.insert(emote.id, emote.code);
+                    }
+                }
+
+                self.emotes.insert(format!("twitch:{}", id), emote_map);
+
                 serde_json::json!({
                     "name": "join_channel",
                     "data": {