From c5517aaaa70ce10cf204fbe02855dc3a487ca961 Mon Sep 17 00:00:00 2001
From: Philippe Weidmann <philippe.weidmann@infomaniak.com>
Date: Thu, 11 May 2023 08:57:02 +0200
Subject: [PATCH] fix: Only set badge for mailboxes with notification enabled

---
 Mail/SceneDelegate.swift                      |  4 +-
 MailCore/Utils/NotificationsHelper.swift      | 39 ++++++++++++-------
 .../NotificationService.swift                 |  6 +--
 3 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/Mail/SceneDelegate.swift b/Mail/SceneDelegate.swift
index de2a6f8f4..f077fcc7f 100644
--- a/Mail/SceneDelegate.swift
+++ b/Mail/SceneDelegate.swift
@@ -60,7 +60,9 @@ class SceneDelegate: UIResponder, UIWindowSceneDelegate, AccountManagerDelegate
     func sceneWillResignActive(_ scene: UIScene) {
         // Called when the scene will move from an active state to an inactive state.
         // This may occur due to temporary interruptions (ex. an incoming phone call).
-        NotificationsHelper.updateUnreadCountBadge()
+        Task {
+            await NotificationsHelper.updateUnreadCountBadge()
+        }
     }
 
     func sceneWillEnterForeground(_ scene: UIScene) {
diff --git a/MailCore/Utils/NotificationsHelper.swift b/MailCore/Utils/NotificationsHelper.swift
index 1dd751171..588de0c42 100644
--- a/MailCore/Utils/NotificationsHelper.swift
+++ b/MailCore/Utils/NotificationsHelper.swift
@@ -21,6 +21,8 @@ import CocoaLumberjackSwift
 import Foundation
 import InfomaniakCore
 import InfomaniakCoreUI
+import InfomaniakDI
+import InfomaniakNotifications
 import MailResources
 import RealmSwift
 import SwiftSoup
@@ -55,17 +57,26 @@ public enum NotificationsHelper {
         }
     }
 
-    public static func getUnreadCount() -> Int {
+    public static func getUnreadCount() async -> Int {
         var totalUnreadCount = 0
-        for mailbox in MailboxInfosManager.instance.getMailboxes() {
-            if let mailboxManager = AccountManager.instance.getMailboxManager(for: mailbox) {
-                totalUnreadCount += mailboxManager.getFolder(with: .inbox)?.unreadCount ?? 0
+        @InjectService var notificationService: InfomaniakNotifications
+
+        for account in AccountManager.instance.accounts {
+            let currentSubscription = await notificationService.subscriptionForUser(id: account.userId)
+
+            for mailbox in MailboxInfosManager.instance.getMailboxes(for: account.userId)
+                where currentSubscription?.topics.contains(mailbox.notificationTopicName) == true {
+                if let mailboxManager = AccountManager.instance.getMailboxManager(for: mailbox) {
+                    totalUnreadCount += mailboxManager.getFolder(with: .inbox)?.unreadCount ?? 0
+                }
             }
         }
+
         return totalUnreadCount
     }
 
-    public static func updateUnreadCountBadge() {
+    @MainActor
+    public static func updateUnreadCountBadge() async {
         // Start a background task to update the app badge when going in the background
         var backgroundTaskIdentifier: UIBackgroundTaskIdentifier = .invalid
         backgroundTaskIdentifier = UIApplication.shared.beginBackgroundTask(withName: "updateUnreadCountBadge task") {
@@ -74,14 +85,12 @@ public enum NotificationsHelper {
             backgroundTaskIdentifier = .invalid
         }
 
-        let totalUnreadCount = getUnreadCount()
+        let totalUnreadCount = await getUnreadCount()
 
-        DispatchQueue.main.async {
-            UIApplication.shared.applicationIconBadgeNumber = totalUnreadCount
-            if backgroundTaskIdentifier != .invalid {
-                UIApplication.shared.endBackgroundTask(backgroundTaskIdentifier)
-                backgroundTaskIdentifier = .invalid
-            }
+        UIApplication.shared.applicationIconBadgeNumber = totalUnreadCount
+        if backgroundTaskIdentifier != .invalid {
+            UIApplication.shared.endBackgroundTask(backgroundTaskIdentifier)
+            backgroundTaskIdentifier = .invalid
         }
     }
 
@@ -114,7 +123,9 @@ public enum NotificationsHelper {
         }
     }
 
-    public static func generateNotificationFor(message: Message, mailboxId: Int, userId: Int) -> UNMutableNotificationContent {
+    public static func generateNotificationFor(message: Message,
+                                               mailboxId: Int,
+                                               userId: Int) async -> UNMutableNotificationContent {
         let content = UNMutableNotificationContent()
         if !message.from.isEmpty {
             content.title = message.from.map { $0.name }.joined(separator: ",")
@@ -125,7 +136,7 @@ public enum NotificationsHelper {
         content.body = getCleanBodyFrom(message: message)
         content.threadIdentifier = "\(mailboxId)_\(userId)"
         content.targetContentIdentifier = "\(userId)_\(mailboxId)_\(message.uid)"
-        content.badge = getUnreadCount() as NSNumber
+        content.badge = await getUnreadCount() as NSNumber
         content.sound = .default
         content.userInfo = [NotificationsHelper.UserInfoKeys.userId: userId,
                             NotificationsHelper.UserInfoKeys.mailboxId: mailboxId,
diff --git a/MailNotificationServiceExtension/NotificationService.swift b/MailNotificationServiceExtension/NotificationService.swift
index 14a8018ac..ef7d558be 100644
--- a/MailNotificationServiceExtension/NotificationService.swift
+++ b/MailNotificationServiceExtension/NotificationService.swift
@@ -108,9 +108,9 @@ class NotificationService: UNNotificationServiceExtension {
                 return contentHandler(bestAttemptContent)
             }
 
-            let completeNotification = NotificationsHelper.generateNotificationFor(message: fetchedMessage,
-                                                                                   mailboxId: mailboxId,
-                                                                                   userId: userId)
+            let completeNotification = await NotificationsHelper.generateNotificationFor(message: fetchedMessage,
+                                                                                         mailboxId: mailboxId,
+                                                                                         userId: userId)
             contentHandler(completeNotification)
         }
     }