Add menu item to auto-assign categories based on nexus assignments

5 files changed, 63 insertions, 0 deletions

diff --git a/src/downloadmanager.cpp b/src/downloadmanager.cpp
index 5ccbdb4d..6cf3a95c 100644
--- a/src/downloadmanager.cpp
+++ b/src/downloadmanager.cpp
@@ -1321,6 +1321,19 @@ QString DownloadManager::getFileName(int index) const
   return m_ActiveDownloads.at(index)->m_FileName;
 }
 
+int DownloadManager::getDownloadIndex(QString filename) const
+{
+  auto file = std::find_if(m_ActiveDownloads.begin(), m_ActiveDownloads.end(), [=](DownloadManager::DownloadInfo *const val) {
+      if (val->m_FileName == filename) return true;
+      return false;
+    });
+  if (file != m_ActiveDownloads.end()) {
+    int fileIndex = m_ActiveDownloads.indexOf(*file);
+    return fileIndex;
+  }
+  return -1;
+}
+
 QDateTime DownloadManager::getFileTime(int index) const
 {
   if ((index < 0) || (index >= m_ActiveDownloads.size())) {
@@ -1389,6 +1402,14 @@ int DownloadManager::getModID(int index) const
   return m_ActiveDownloads.at(index)->m_FileInfo->modID;
 }
 
+int DownloadManager::getCategoryID(int index) const
+{
+  if ((index < 0) || (index >= m_ActiveDownloads.size())) {
+    throw MyException(tr("mod id: invalid download index %1").arg(index));
+  }
+  return m_ActiveDownloads.at(index)->m_FileInfo->categoryID;
+}
+
 QString DownloadManager::getDisplayGameName(int index) const
 {
   if ((index < 0) || (index >= m_ActiveDownloads.size())) {
diff --git a/src/downloadmanager.h b/src/downloadmanager.h
index 305c10e3..359455d9 100644
--- a/src/downloadmanager.h
+++ b/src/downloadmanager.h
@@ -299,6 +299,14 @@ public:
   QString getFileName(int index) const;
 
   /**
+   * @brief retrieve the file index from the filename
+   *
+   * @param filename the filename of the download
+   * @return the index of the file
+   */
+  int getDownloadIndex(QString filename) const;
+
+  /**
    * @brief retrieve the file size of the download specified by index
    *
    * @param index index of the file to look up
@@ -349,6 +357,14 @@ public:
   int getModID(int index) const;
 
   /**
+   * @brief retrieve the nexus category id of the download specified by index
+   *
+   * @param index index of the file to look up
+   * @return the nexus category id
+   */
+  int getCategoryID(int index) const;
+
+  /**
    * @brief retrieve the displayable game name of the download specified by the index
    *
    * @param index index of the file to look up
diff --git a/src/modlistcontextmenu.cpp b/src/modlistcontextmenu.cpp
index 8c9c0206..66766b0e 100644
--- a/src/modlistcontextmenu.cpp
+++ b/src/modlistcontextmenu.cpp
@@ -94,6 +94,9 @@ void ModListGlobalContextMenu::populate(OrganizerCore& core, ModListView* view,
   addAction(tr("Check for updates"), [=]() {
     view->actions().checkModsForUpdates();
   });
+  addAction(tr("Auto assign categories"), [=]() {
+   view->actions().assignCategories();
+  });
   addAction(tr("Refresh"), &core, &OrganizerCore::profileRefresh);
   addAction(tr("Export to csv..."), [=]() {
     view->actions().exportModListCSV();
diff --git a/src/modlistviewactions.cpp b/src/modlistviewactions.cpp
index a6f0f07e..21079f1d 100644
--- a/src/modlistviewactions.cpp
+++ b/src/modlistviewactions.cpp
@@ -12,6 +12,7 @@
 #include "categories.h"
 #include "csvbuilder.h"
 #include "directoryrefresher.h"
+#include "downloadmanager.h"
 #include "filedialogmemory.h"
 #include "filterlist.h"
 #include "listdialog.h"
@@ -259,6 +260,23 @@ void ModListViewActions::checkModsForUpdates() const
   }
 }
 
+void ModListViewActions::assignCategories() const
+{
+  for (auto mod : m_core.modList()->allMods()) {
+    ModInfo::Ptr modInfo = ModInfo::getByName(mod);
+    for (auto category : modInfo->categories()) {
+      modInfo->removeCategory(category);
+    }
+    QString file = modInfo->installationFile();
+    auto download = m_core.downloadManager()->getDownloadIndex(file);
+    if (download >= 0) {
+      int nexusCategory = m_core.downloadManager()->getCategoryID(download);
+      int category = CategoryFactory::instance()->resolveNexusID(nexusCategory);
+      modInfo->setCategory(CategoryFactory::instance()->getCategoryID(category), true);
+    }
+  }
+}
+
 void ModListViewActions::checkModsForUpdates(
     std::multimap<QString, int> const& IDs) const
 {
diff --git a/src/modlistviewactions.h b/src/modlistviewactions.h
index 5927654f..3805f98b 100644
--- a/src/modlistviewactions.h
+++ b/src/modlistviewactions.h
@@ -13,6 +13,7 @@ class MainWindow;
 class ModListView;
 class PluginListView;
 class OrganizerCore;
+class DownloadManager;
 
 class ModListViewActions : public QObject
 {
@@ -53,6 +54,10 @@ public:
   void checkModsForUpdates() const;
   void checkModsForUpdates(const QModelIndexList& indices) const;
 
+  // auto-assign categories based on nexus ID
+  //
+  void assignCategories() const;
+
   // start the "Export Mod List" dialog
   //
   void exportModListCSV() const;