From d2300bf663a797f40849b5ec8914659d7263f18d Mon Sep 17 00:00:00 2001
From: miruka <miruka@disroot.org>
Date: Sun, 31 May 2020 19:41:05 -0400
Subject: [PATCH] Update account total notif/highlight calculation

---
 src/backend/matrix_client.py | 33 ++++++++++++++++++++++++++++-----
 1 file changed, 28 insertions(+), 5 deletions(-)

diff --git a/src/backend/matrix_client.py b/src/backend/matrix_client.py
index 0c0a9fe8..3502641b 100644
--- a/src/backend/matrix_client.py
+++ b/src/backend/matrix_client.py
@@ -1212,6 +1212,21 @@ class MatrixClient(nio.AsyncClient):
         self.models[self.user_id, "rooms"][room_id].last_event_date = \
             ZeroDate
 
+
+    async def update_account_unread_counts(self) -> None:
+        """Recalculate total unread notifications/highlights for our account"""
+
+        unreads = highlights = 0
+
+        for room in self.models[self.user_id, "rooms"].values():
+            unreads    += room.unreads
+            highlights += room.highlights
+
+        account                  = self.models["accounts"][self.user_id]
+        account.total_unread     = unreads
+        account.total_highlights = highlights
+
+
     # Functions to register data into models
 
     async def event_is_past(self, ev: Union[nio.Event, Event]) -> bool:
@@ -1256,12 +1271,17 @@ class MatrixClient(nio.AsyncClient):
         try:
             registered = self.models[self.user_id, "rooms"][room.room_id]
         except KeyError:
-            registered      = None
-            last_event_date = datetime.fromtimestamp(0)
-            typing_members  = []
+            registered                   = None
+            last_event_date              = datetime.fromtimestamp(0)
+            typing_members               = []
+            update_account_unread_counts = True
         else:
-            last_event_date = registered.last_event_date
-            typing_members  = registered.typing_members
+            last_event_date              = registered.last_event_date
+            typing_members               = registered.typing_members
+            update_account_unread_counts = (
+                registered.unreads != room.unread_notifications or
+                registered.highlights != room.unread_highlights
+            )
 
         room_item = Room(
             id             = room.room_id,
@@ -1308,6 +1328,9 @@ class MatrixClient(nio.AsyncClient):
             for user_id in room.users:
                 await self.add_member(room, user_id)
 
+        if update_account_unread_counts:
+            await self.update_account_unread_counts()
+
 
     async def add_member(self, room: nio.MatrixRoom, user_id: str) -> None:
         """Register/update a room member into our models."""