Refactor and added in NotificationService