4 files changed, 55 insertions, 7 deletions

diff --git a/src/modlistcontextmenu.cpp b/src/modlistcontextmenu.cpp
index 01c4b471..19b5fba4 100644
--- a/src/modlistcontextmenu.cpp
+++ b/src/modlistcontextmenu.cpp
@@ -126,13 +126,9 @@ void ModListPrimaryCategoryMenu::populate(const CategoryFactory& factory, ModInf
       QRadioButton* categoryBox = new QRadioButton(
         factory.getCategoryName(catIdx).replace('&', "&&"),
         this);
-      connect(categoryBox, &QRadioButton::toggled, [mod, categoryID](bool enable) {
-        if (enable) {
-          mod->setPrimaryCategory(categoryID);
-        }
-      });
       categoryBox->setChecked(categoryID == mod->primaryCategory());
       action->setDefaultWidget(categoryBox);
+      action->setData(categoryID);
     }
     catch (const std::exception& e) {
       log::error("failed to create category checkbox: {}", e.what());
@@ -143,6 +139,20 @@ void ModListPrimaryCategoryMenu::populate(const CategoryFactory& factory, ModInf
   }
 }
 
+int ModListPrimaryCategoryMenu::primaryCategory() const
+{
+  for (QAction* action : actions()) {
+    QWidgetAction* widgetAction = qobject_cast<QWidgetAction*>(action);
+    if (widgetAction) {
+      QRadioButton* button = qobject_cast<QRadioButton*>(widgetAction->defaultWidget());
+      if (button && button->isChecked()) {
+        return widgetAction->data().toInt();
+      }
+    }
+  }
+  return -1;
+}
+
 ModListContextMenu::ModListContextMenu(
   const QModelIndex& index, OrganizerCore& core, CategoryFactory& categories, ModListView* view) :
   QMenu(view)
@@ -245,6 +255,12 @@ void ModListContextMenu::addSeparatorActions(ModInfo::Ptr mod)
   addMenuAsPushButton(categoriesMenu);
 
   ModListPrimaryCategoryMenu* primaryCategoryMenu = new ModListPrimaryCategoryMenu(m_categories, mod, this);
+  connect(primaryCategoryMenu, &QMenu::aboutToHide, [=]() {
+    int category = primaryCategoryMenu->primaryCategory();
+    if (category != -1) {
+      m_actions.setPrimaryCategory(m_selected, category);
+    }
+  });
   addMenuAsPushButton(primaryCategoryMenu);
   addSeparator();
 
@@ -310,6 +326,12 @@ void ModListContextMenu::addRegularActions(ModInfo::Ptr mod)
   addMenuAsPushButton(categoriesMenu);
 
   ModListPrimaryCategoryMenu* primaryCategoryMenu = new ModListPrimaryCategoryMenu(m_categories, mod, this);
+  connect(primaryCategoryMenu, &QMenu::aboutToHide, [=]() {
+    int category = primaryCategoryMenu->primaryCategory();
+    if (category != -1) {
+      m_actions.setPrimaryCategory(m_selected, category);
+    }
+  });
   addMenuAsPushButton(primaryCategoryMenu);
   addSeparator();
 
diff --git a/src/modlistcontextmenu.h b/src/modlistcontextmenu.h
index d009c9d8..aa9d4c2b 100644
--- a/src/modlistcontextmenu.h
+++ b/src/modlistcontextmenu.h
@@ -55,6 +55,10 @@ public:
 
   ModListPrimaryCategoryMenu(CategoryFactory& categories, ModInfo::Ptr mod, QMenu* parent = nullptr);
 
+  // return the selected primary category
+  //
+  int primaryCategory() const;
+
 private:
 
   // populate the categories
diff --git a/src/modlistviewactions.cpp b/src/modlistviewactions.cpp
index 8d24a8d5..cb212e0b 100644
--- a/src/modlistviewactions.cpp
+++ b/src/modlistviewactions.cpp
@@ -956,7 +956,6 @@ void ModListViewActions::setCategories(const QModelIndexList& selected, const QM
 {
   ModInfo::Ptr refMod = ModInfo::getByIndex(ref.data(ModList::IndexRole).toInt());
   if (selected.size() > 1) {
-
     for (auto& idx : selected) {
       if (idx.row() != ref.row()) {
         setCategoriesIf(
@@ -982,6 +981,23 @@ void ModListViewActions::setCategories(const QModelIndexList& selected, const QM
   }
 }
 
+void ModListViewActions::setPrimaryCategory(const QModelIndexList& selected, int category, bool force)
+{
+  for (auto& idx : selected) {
+    ModInfo::Ptr info = ModInfo::getByIndex(idx.data(ModList::IndexRole).toInt());
+    if (force || info->categorySet(category)) {
+      info->setCategory(category, true);
+      info->setPrimaryCategory(category);
+    }
+  }
+
+  // reset the selection manually - still needed
+  auto viewIndices = m_view->indexModelToView(selected);
+  for (auto& idx : viewIndices) {
+    m_view->selectionModel()->select(idx, QItemSelectionModel::Select | QItemSelectionModel::Rows);
+  }
+}
+
 void ModListViewActions::openExplorer(const QModelIndexList& index) const
 {
   for (auto& idx : index) {
diff --git a/src/modlistviewactions.h b/src/modlistviewactions.h
index 65f323d9..2850d30a 100644
--- a/src/modlistviewactions.h
+++ b/src/modlistviewactions.h
@@ -87,7 +87,7 @@ public:
   void setColor(const QModelIndexList& indices, const QModelIndex& refIndex) const;
   void resetColor(const QModelIndexList& indices) const;
 
-  // set the category of the mod in the given list, using the given index as reference
+  // set the category of the mods in the given list, using the given index as reference
   // - the categories are set as-is on the refernce mod
   // - for the other mods, the category is only set if the current state of the category
   //   on the reference is different
@@ -95,6 +95,12 @@ public:
   void setCategories(const QModelIndexList& selected, const QModelIndex& ref,
     const std::vector<std::pair<int, bool>>& categories) const;
 
+  // set the primary category of the mods in the given list, adding the appropriate
+  // category to the mods unless force is false, in which case the primary category
+  // is set only on mods that have this category
+  //
+  void setPrimaryCategory(const QModelIndexList& selected, int category, bool force = true);
+
   // open the Windows explorer for the specified mods
   //
   void openExplorer(const QModelIndexList& index) const;