4 files changed, 48 insertions, 20 deletions

diff --git a/src/system/bot/client.rs b/src/system/bot/client.rs
index 4d4f7bb..a2d1b28 100644
--- a/src/system/bot/client.rs
+++ b/src/system/bot/client.rs
@@ -35,6 +35,18 @@ impl Client {
             .expect("Could not deserialize message")
     }
 
+    pub async fn fetch_recent_channel_messages(&self, channel_id: ChannelId) -> Result<Vec<FullMessage>, TwiError> {
+        let client = self.client.lock().await;
+
+        Ok(client
+            .channel_messages(channel_id)
+            .limit(10).unwrap()
+            .await?
+            .model()
+            .await
+            .unwrap())
+    }
+
     pub async fn resend_message(&self, message_id: MessageId, channel_id: ChannelId) {
         let bot_conf = self.bot_conf.read().await;
         let message = self.fetch_message(message_id, channel_id).await;
diff --git a/src/system/bot/mod.rs b/src/system/bot/mod.rs
index f0a2e45..2aa3e0a 100644
--- a/src/system/bot/mod.rs
+++ b/src/system/bot/mod.rs
@@ -70,6 +70,10 @@ impl Bot {
         self.client.fetch_message(message_id, channel_id).await
     }
 
+    pub async fn fetch_recent_channel_messages(&self, channel_id: ChannelId) -> Result<Vec<FullMessage>, TwiError> {
+        self.client.fetch_recent_channel_messages(channel_id).await
+    }
+
     pub async fn resend_message(&self, message_id: MessageId, channel_id: ChannelId) {
         self.client.resend_message(message_id, channel_id).await;
     }
diff --git a/src/system/message_parser.rs b/src/system/message_parser.rs
index d161483..b044f61 100644
--- a/src/system/message_parser.rs
+++ b/src/system/message_parser.rs
@@ -39,7 +39,7 @@ static CORRECTION_REGEX: LazyLock<Regex> = LazyLock::new(|| {
 });
 
 impl MessageParser {
-    pub fn parse(message: &FullMessage, secondary_message: Option<FullMessage>, system_config: &System, latch_state: Option<(MemberId, Timestamp)>) -> ParsedMessage {
+    pub fn parse(message: &FullMessage, secondary_message: Option<&FullMessage>, system_config: &System, latch_state: Option<(MemberId, Timestamp)>) -> ParsedMessage {
         if message.content == r"\\" {
             return ParsedMessage::LatchClear(if let Some((member_id, _)) = latch_state {
                 member_id
@@ -76,7 +76,7 @@ impl MessageParser {
         ParsedMessage::UnproxiedMessage
     }
 
-    fn parse_command(message: &FullMessage, secondary_message: Option<FullMessage>, system_config: &System, latch_state: Option<(MemberId, Timestamp)>) -> Command {
+    fn parse_command(message: &FullMessage, secondary_message: Option<&FullMessage>, system_config: &System, latch_state: Option<(MemberId, Timestamp)>) -> Command {
         let mut words = message.content.strip_prefix("!").unwrap().split_whitespace();
         let first_word = words.next();
 
@@ -145,7 +145,7 @@ impl MessageParser {
         Command::UnknownCommand
     }
 
-    fn check_correction(message: &FullMessage, secondary_message: Option<FullMessage>) -> Option<ParsedMessage> {
+    fn check_correction(message: &FullMessage, secondary_message: Option<&FullMessage>) -> Option<ParsedMessage> {
         None
     }
 
diff --git a/src/system/mod.rs b/src/system/mod.rs
index b5cd1ab..6d81536 100644
--- a/src/system/mod.rs
+++ b/src/system/mod.rs
@@ -167,27 +167,39 @@ impl Manager {
     }
 
     async fn handle_message(&mut self, message: TwiMessage, timestamp: Timestamp, seen_by: MemberId) {
-        // let bot = self.bots.get(&seen_by).expect("No client for member");
-        let last_in_channel = self.send_cache.get(&message.channel_id);
-        let replied_message = if let MessageType::Reply = message.kind {
-            message.referenced_message.clone()
-        } else {
-            None
-        };
-
-        if let None = last_in_channel {
-            println!("ERROR: Could not look up last sent message in channel {}", message.channel_id);
-        }
+        let bot = self.bots.get(&seen_by).expect("No client for member");
 
-        let ref_message = if replied_message.is_some() {
-            replied_message.map(|m| *m)
-        } else if last_in_channel.is_some() {
-            last_in_channel.map(|m| m.clone())
+        // If message type is reply, use that
+        let referenced_message = if let MessageType::Reply = message.kind {
+            message.referenced_message.as_ref().map(|message| message.as_ref())
         } else {
-            None
+            // Otherwise, check cache for lest message sent in channel
+            if self.send_cache.contains(&message.channel_id) {
+                self.send_cache.get(&message.channel_id)
+            } else {
+                // Or look it up if it's not in cache
+                let system_bot_ids : Vec<UserId> = self.config.members.iter().filter_map(|m| m.user_id).collect();
+                let recent_messages = bot.fetch_recent_channel_messages(message.channel_id).await;
+
+                let last_in_channel = recent_messages.map(|messages| {
+                    messages.into_iter().filter(|message|
+                        system_bot_ids.contains(&message.author.id)
+                    ).max_by_key(|message| message.timestamp.as_micros())
+                }).ok().flatten();
+
+                // Since we did all this work to look it up, insert it into cache
+                if let Some(last) = last_in_channel {
+                    self.send_cache.put(message.channel_id, last);
+                } else {
+                    println!("WARNING: Could not look up most recent message in channel {}", message.channel_id);
+                };
+
+                // Return the message referenced from cache so there's no unnecessary clone
+                self.send_cache.get(&message.channel_id)
+            }
         };
 
-        let parsed_message = MessageParser::parse(&message, ref_message, &self.config, self.latch_state);
+        let parsed_message = MessageParser::parse(&message, referenced_message, &self.config, self.latch_state);
 
         match parsed_message {
             message_parser::ParsedMessage::UnproxiedMessage => (),