From 3e2888eb75268257e7c5e0425f5f0e90d4f8604e Mon Sep 17 00:00:00 2001
From: Chewbacca <dice.warrant-02@icloud.com>
Date: Thu, 27 Oct 2022 13:59:54 -0400
Subject: [PATCH] Handle unreadCount when messages are deleted

---
 app/soapbox/actions/streaming.ts |  6 +++---
 app/soapbox/queries/chats.ts     | 21 +++++++++++++++++----
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/app/soapbox/actions/streaming.ts b/app/soapbox/actions/streaming.ts
index 4b4872bc0..3f671967d 100644
--- a/app/soapbox/actions/streaming.ts
+++ b/app/soapbox/actions/streaming.ts
@@ -78,10 +78,10 @@ const updateChat = (payload: ChatPayload) => {
 };
 
 const removeChatMessage = (payload: string) => {
-  const chat = JSON.parse(payload);
-  const chatMessageId = chat.chat_message_id;
+  const data = JSON.parse(payload);
+  const chatMessageId = data.deleted_message_id;
 
-  removePageItem(ChatKeys.chatMessages(chat.id), chatMessageId, (o: any, n: any) => Number(o.id) === Number(n));
+  removePageItem(ChatKeys.chatMessages(data.chat_id), chatMessageId, (o: any, n: any) => String(o.id) === String(n));
 };
 
 const connectTimelineStream = (
diff --git a/app/soapbox/queries/chats.ts b/app/soapbox/queries/chats.ts
index c22f0ff7b..b6fa3e66f 100644
--- a/app/soapbox/queries/chats.ts
+++ b/app/soapbox/queries/chats.ts
@@ -1,4 +1,4 @@
-import { useInfiniteQuery, useMutation, useQuery } from '@tanstack/react-query';
+import { InfiniteData, useInfiniteQuery, useMutation, useQuery } from '@tanstack/react-query';
 import sumBy from 'lodash/sumBy';
 
 import { fetchRelationships } from 'soapbox/actions/accounts';
@@ -130,8 +130,7 @@ const useChats = (search?: string) => {
     const link = getNextLink(response);
     const hasMore = !!link;
 
-    // TODO: change to response header
-    setUnreadChatsCount(sumBy(data, (chat) => chat.unread));
+    setUnreadChatsCount(Number(response.headers['x-unread-messages-count']) || sumBy(data, (chat) => chat.unread));
 
     // Set the relationships to these users in the redux store.
     dispatch(fetchRelationships(data.map((item) => item.account.id)));
@@ -190,12 +189,26 @@ const useChat = (chatId?: string) => {
 const useChatActions = (chatId: string) => {
   const api = useApi();
   const dispatch = useAppDispatch();
+  const { setUnreadChatsCount } = useStatContext();
 
   const { chat, setChat, setEditing } = useChatContext();
 
   const markChatAsRead = (lastReadId: string) => {
     api.post<IChat>(`/api/v1/pleroma/chats/${chatId}/read`, { last_read_id: lastReadId })
-      .then(({ data }) => updatePageItem(['chats', 'search'], data, (o, n) => o.id === n.id))
+      .then(({ data }) => {
+        updatePageItem(ChatKeys.chatSearch(), data, (o, n) => o.id === n.id);
+        const queryData = queryClient.getQueryData<InfiniteData<PaginatedResult<unknown>>>(ChatKeys.chatSearch());
+        if (queryData) {
+          const flattenedQueryData: any = flattenPages(queryData)?.map((chat: any) => {
+            if (chat.id === data.id) {
+              return data;
+            } else {
+              return chat;
+            }
+          });
+          setUnreadChatsCount(sumBy(flattenedQueryData, (chat: IChat) => chat.unread));
+        }
+      })
       .catch(() => null);
   };