Add private message notification as a badge to the menu instead of using a block

I have a menu item in my main navigation linking to the inbox. Instead of using the notification block I would like to add a badge to the message menu link item. I would like to use as much of your code so I tried the following which works fine.

However I am wondering whether there is a better way of doing it as I think its tricky to add the needed classes.

Also maybe the following code can be usefull for someone or added to the module in this or a modified way. I can imagine that other users also might prefer using a badge on an existing menu item rather than a new block. In this case the refresh rate should be implemented as a general setting instead of beeing fixed like in my code.

The code lacks of handling caching properly yet. The badge is correctly refreshed by your js, but when reloading the page it shows a wrong value. I don't know whether it makes sense to disable caching for the menu only for the badge. Maybe a better approach would be to also handle the initial setting of the value on initialization of the JS file.

function my_module_preprocess_menu(&$variables) {
  
  if(isset($variables['menu_name']) && $variables['menu_name'] == 'main') {
    
    $items = &$variables['items'];
    
    foreach ($items as &$item) {
      
      if ($item['title'] == 'Messages') {
        
        $url = $item['url'];
        $link_options = array(
          'attributes' => array(
            'class' => array(
              'private-message-notification-wrapper',
            ),
          ),
        );
        
        $url->setOptions($link_options);
        $item['url'] = $url;
        
        $item['title'] = new FormattableMarkup('@title <span class="private-message-page-link badge">@count</span>', [
            '@title' => $item['title'],
            '@count' => \Drupal::service('private_message.service')->getUnreadThreadCount(),
          ]);
          
      }
      
    }
    
    $url = Url::fromRoute('private_message.ajax_callback', ['op' => 'get_new_unread_thread_count']);
    $token = \Drupal::service('csrf_token')->get($url->getInternalPath());
    $url->setOptions(['absolute' => TRUE, 'query' => ['token' => $token]]);
    $variables['#attached']['drupalSettings']['privateMessageNotificationBlock']['newMessageCountCallback'] = $url->toString();
    $variables['#attached']['drupalSettings']['privateMessageNotificationBlock']['ajaxRefreshRate'] = 60;

    $variables['#attached']['library'][]  = 'private_message/notification_block';
    
  }
}