2 files changed, 33 insertions, 24 deletions

diff --git a/launcher/ui/dialogs/ResourceDownloadDialog.cpp b/launcher/ui/dialogs/ResourceDownloadDialog.cpp
index 523a1636..2eb85928 100644
--- a/launcher/ui/dialogs/ResourceDownloadDialog.cpp
+++ b/launcher/ui/dialogs/ResourceDownloadDialog.cpp
@@ -141,38 +141,44 @@ ResourcePage* ResourceDownloadDialog::getSelectedPage()
     return m_selectedPage;
 }
 
-void ResourceDownloadDialog::addResource(QString name, ResourceDownloadTask* task)
+void ResourceDownloadDialog::addResource(ModPlatform::IndexedPack& pack, ModPlatform::IndexedVersion& ver, bool is_indexed)
 {
-    removeResource(name);
-    m_selected.insert(name, task);
+    removeResource(pack, ver);
+
+    ver.is_currently_selected = true;
+    m_selected.insert(pack.name, new ResourceDownloadTask(pack, ver, getBaseModel(), is_indexed));
 
     m_buttons.button(QDialogButtonBox::Ok)->setEnabled(!m_selected.isEmpty());
 }
 
-void ResourceDownloadDialog::removeResource(QString name)
+static ModPlatform::IndexedVersion& getVersionWithID(ModPlatform::IndexedPack& pack, QVariant id)
 {
-    if (m_selected.contains(name))
-        m_selected.find(name).value()->deleteLater();
-    m_selected.remove(name);
-
-    m_buttons.button(QDialogButtonBox::Ok)->setEnabled(!m_selected.isEmpty());
+    Q_ASSERT(pack.versionsLoaded);
+    auto it = std::find_if(pack.versions.begin(), pack.versions.end(), [id](auto const& v) { return v.fileId == id; });
+    Q_ASSERT(it != pack.versions.end());
+    return *it;
 }
 
-bool ResourceDownloadDialog::isSelected(QString name, QString filename) const
+void ResourceDownloadDialog::removeResource(ModPlatform::IndexedPack& pack, ModPlatform::IndexedVersion& ver)
 {
-    auto iter = m_selected.constFind(name);
-    if (iter == m_selected.constEnd())
-        return false;
+    if (auto selected_task_it = m_selected.find(pack.name); selected_task_it != m_selected.end()) {
+        auto selected_task = *selected_task_it;
+        auto old_version_id = selected_task->getVersionID();
 
-    // FIXME: Is there a way to check for versions without checking the filename
-    //        as a heuristic, other than adding such info to ResourceDownloadTask itself?
-    if (!filename.isEmpty())
-        return iter.value()->getFilename() == filename;
+        // If the new and old version IDs don't match, search for the old one and deselect it.
+        if (ver.fileId != old_version_id)
+            getVersionWithID(pack, old_version_id).is_currently_selected = false;
+    }
 
-    return true;
+    // Deselect the new version too, since all versions of that pack got removed.
+    ver.is_currently_selected = false;
+
+    m_selected.remove(pack.name);
+
+    m_buttons.button(QDialogButtonBox::Ok)->setEnabled(!m_selected.isEmpty());
 }
 
-const QList<ResourceDownloadTask*> ResourceDownloadDialog::getTasks()
+const QList<ResourceDownloadDialog::DownloadTaskPtr> ResourceDownloadDialog::getTasks()
 {
     return m_selected.values();
 }
diff --git a/launcher/ui/dialogs/ResourceDownloadDialog.h b/launcher/ui/dialogs/ResourceDownloadDialog.h
index 29813493..95a5e628 100644
--- a/launcher/ui/dialogs/ResourceDownloadDialog.h
+++ b/launcher/ui/dialogs/ResourceDownloadDialog.h
@@ -23,6 +23,8 @@
 #include <QDialogButtonBox>
 #include <QLayout>
 
+#include "QObjectPtr.h"
+#include "modplatform/ModIndex.h"
 #include "ui/pages/BasePageProvider.h"
 
 class BaseInstance;
@@ -41,6 +43,8 @@ class ResourceDownloadDialog : public QDialog, public BasePageProvider {
     Q_OBJECT
 
    public:
+    using DownloadTaskPtr = shared_qobject_ptr<ResourceDownloadTask>;
+
     ResourceDownloadDialog(QWidget* parent, const std::shared_ptr<ResourceFolderModel> base_model);
 
     void initializeContainer();
@@ -54,11 +58,10 @@ class ResourceDownloadDialog : public QDialog, public BasePageProvider {
     bool selectPage(QString pageId);
     ResourcePage* getSelectedPage();
 
-    void addResource(QString name, ResourceDownloadTask* task);
-    void removeResource(QString name);
-    [[nodiscard]] bool isSelected(QString name, QString filename = "") const;
+    void addResource(ModPlatform::IndexedPack&, ModPlatform::IndexedVersion&, bool is_indexed = false);
+    void removeResource(ModPlatform::IndexedPack&, ModPlatform::IndexedVersion&);
 
-    const QList<ResourceDownloadTask*> getTasks();
+    const QList<DownloadTaskPtr> getTasks();
     [[nodiscard]] const std::shared_ptr<ResourceFolderModel> getBaseModel() const { return m_base_model; }
 
    public slots:
@@ -82,7 +85,7 @@ class ResourceDownloadDialog : public QDialog, public BasePageProvider {
     QDialogButtonBox m_buttons;
     QVBoxLayout m_vertical_layout;
 
-    QHash<QString, ResourceDownloadTask*> m_selected;
+    QHash<QString, DownloadTaskPtr> m_selected;
 };