From 7e9646cc25e09aab594f10d6a2210b550d9fb0af Mon Sep 17 00:00:00 2001
From: Loren Burkholder <computersemiexpert@outlook.com>
Date: Mon, 25 Apr 2022 20:54:40 -0400
Subject: [PATCH] Add notifications to all sidebar items

---
 src/timeline/CommunitiesModel.cpp | 53 +++++++++++++++++++++++++------
 src/timeline/CommunitiesModel.h   |  2 ++
 2 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/src/timeline/CommunitiesModel.cpp b/src/timeline/CommunitiesModel.cpp
index c89e7661..34da6157 100644
--- a/src/timeline/CommunitiesModel.cpp
+++ b/src/timeline/CommunitiesModel.cpp
@@ -20,7 +20,7 @@ CommunitiesModel::CommunitiesModel(QObject *parent)
     connect(ChatPage::instance(), &ChatPage::unreadMessages, this, [this](int) {
         // Simply updating every space is easier than tracking which ones need updated.
         if (!spaces_.empty())
-            emit dataChanged(index(2, 0), index(spaces_.size() + 2, 0), {Roles::UnreadMessages, Roles::HasLoudNotification});
+            emit dataChanged(index(0, 0), index(spaces_.size() + tags_.size() + 1, 0), {Roles::UnreadMessages, Roles::HasLoudNotification});
     });
 }
 
@@ -80,9 +80,17 @@ CommunitiesModel::data(const QModelIndex &index, int role) const
             return 0;
         case CommunitiesModel::Roles::Id:
             return "";
-        case CommunitiesModel::Roles::UnreadMessages:
+        case CommunitiesModel::Roles::UnreadMessages: {
+            int total{0};
+            for (const auto &[id, info] : cache::getRoomInfo(cache::joinedRooms()))
+                total += info.notification_count;
+            return total;
+        }
         case CommunitiesModel::Roles::HasLoudNotification:
-            return 0;
+            for (const auto &[id, info] : cache::getRoomInfo(cache::joinedRooms()))
+                if (info.highlight_count > 0)
+                    return true;
+            return false;
         }
     } else if (index.row() == 1) {
         switch (role) {
@@ -104,9 +112,17 @@ CommunitiesModel::data(const QModelIndex &index, int role) const
             return 0;
         case CommunitiesModel::Roles::Id:
             return "dm";
-        case CommunitiesModel::Roles::UnreadMessages:
+        case CommunitiesModel::Roles::UnreadMessages: {
+            int total{0};
+            for (const auto &[id, info] : cache::getRoomInfo(directMessages_))
+                total += info.notification_count;
+            return total;
+        }
         case CommunitiesModel::Roles::HasLoudNotification:
-            return 0;
+            for (const auto &[id, info] : cache::getRoomInfo(directMessages_))
+                if (info.highlight_count > 0)
+                    return true;
+            return false;
         }
     } else if (index.row() - 2 < spaceOrder_.size()) {
         auto id = spaceOrder_.tree.at(index.row() - 2).id;
@@ -191,9 +207,22 @@ CommunitiesModel::data(const QModelIndex &index, int role) const
             return 0;
         case CommunitiesModel::Roles::Id:
             return "tag:" + tag;
-        case CommunitiesModel::Roles::UnreadMessages:
-        case CommunitiesModel::Roles::HasLoudNotification:
-            return 0;
+        case CommunitiesModel::Roles::UnreadMessages: {
+            int total{0};
+            auto rooms{cache::joinedRooms()};
+            for (const auto &[roomid, info] : cache::getRoomInfo(rooms))
+                if (std::find(std::begin(info.tags), std::end(info.tags), tag.toStdString()) != std::end(info.tags))
+                    total += info.notification_count;
+            return total;
+        }
+        case CommunitiesModel::Roles::HasLoudNotification: {
+            auto rooms{cache::joinedRooms()};
+            for (const auto &[roomid, info] : cache::getRoomInfo(rooms))
+                if (std::find(std::begin(info.tags), std::end(info.tags), tag.toStdString()) != std::end(info.tags))
+                    if (info.highlight_count > 0)
+                        return true;
+            return false;
+        }
         }
     }
     return QVariant();
@@ -403,8 +432,12 @@ CommunitiesModel::sync(const mtx::responses::Sync &sync_)
             tagsUpdated = true;
     }
     for (const auto &e : sync_.account_data.events) {
-        if (std::holds_alternative<
-              mtx::events::AccountDataEvent<mtx::events::account_data::Direct>>(e)) {
+        if (auto event =
+              std::get_if<mtx::events::AccountDataEvent<mtx::events::account_data::Direct>>(&e)) {
+            directMessages_.clear();
+            for (const auto &[userId, roomIds] : event->content.user_to_rooms)
+                for (const auto &roomId : roomIds)
+                    directMessages_.push_back(roomId);
             tagsUpdated = true;
             break;
         }
diff --git a/src/timeline/CommunitiesModel.h b/src/timeline/CommunitiesModel.h
index 7be98939..bcb487e7 100644
--- a/src/timeline/CommunitiesModel.h
+++ b/src/timeline/CommunitiesModel.h
@@ -50,6 +50,7 @@ public:
         Id,
         UnreadMessages,
         HasLoudNotification,
+        IsDirect,
     };
 
     struct FlatTree
@@ -154,6 +155,7 @@ private:
     QStringList hiddenTagIds_;
     FlatTree spaceOrder_;
     std::map<QString, RoomInfo> spaces_;
+    std::vector<std::string> directMessages_;
 
     friend class FilteredCommunitiesModel;
 };