From 588c4aa0bddc8ae80d64baebfa442fa5c76bcb82 Mon Sep 17 00:00:00 2001
From: Aruki <arukibree@github.com>
Date: Mon, 24 Jul 2017 23:51:09 -0600
Subject: [PATCH] Added support for building Wii de Asobu/Trilogy ISOs

---
 src/Core/GameProject/CGameProject.cpp | 20 ++++++++++
 src/Core/GameProject/CGameProject.h   |  6 ++-
 src/Editor/CProjectSettingsDialog.cpp | 53 ++++++++++++++++++++++++++-
 3 files changed, 76 insertions(+), 3 deletions(-)

diff --git a/src/Core/GameProject/CGameProject.cpp b/src/Core/GameProject/CGameProject.cpp
index e0e8217f..2f0fa61c 100644
--- a/src/Core/GameProject/CGameProject.cpp
+++ b/src/Core/GameProject/CGameProject.cpp
@@ -80,6 +80,7 @@ bool CGameProject::Serialize(IArchive& rArc)
 bool CGameProject::BuildISO(const TString& rkIsoPath, IProgressNotifier *pProgress)
 {
     ASSERT( FileUtil::IsValidPath(rkIsoPath, false) );
+    ASSERT( !IsWiiDeAsobu() && !IsTrilogy() );
 
     auto ProgressCallback = [&](float ProgressPercent, const nod::SystemString& rkInfoString, size_t)
     {
@@ -101,6 +102,25 @@ bool CGameProject::BuildISO(const TString& rkIsoPath, IProgressNotifier *pProgre
     }
 }
 
+bool CGameProject::MergeISO(const TString& rkIsoPath, nod::DiscWii *pOriginalIso, IProgressNotifier *pProgress)
+{
+    ASSERT( FileUtil::IsValidPath(rkIsoPath, false) );
+    ASSERT( IsWiiDeAsobu() || IsTrilogy() );
+    ASSERT( pOriginalIso != nullptr );
+
+    auto ProgressCallback = [&](float ProgressPercent, const nod::SystemString& rkInfoString, size_t)
+    {
+        pProgress->Report((int) (ProgressPercent * 10000), 10000, TWideString(rkInfoString).ToUTF8());
+    };
+
+    pProgress->SetTask(0, "Building " + rkIsoPath.GetFileName());
+
+    TWideString DiscRoot = DiscFilesystemRoot(false).ToUTF16();
+
+    nod::DiscMergerWii Merger(*rkIsoPath.ToUTF16(), *pOriginalIso, IsTrilogy(), ProgressCallback);
+    return Merger.mergeFromDirectory(*DiscRoot) == nod::EBuildResult::Success;
+}
+
 void CGameProject::GetWorldList(std::list<CAssetID>& rOut) const
 {
     for (u32 iPkg = 0; iPkg < mPackages.size(); iPkg++)
diff --git a/src/Core/GameProject/CGameProject.h b/src/Core/GameProject/CGameProject.h
index fdd86d7b..19a0c2ab 100644
--- a/src/Core/GameProject/CGameProject.h
+++ b/src/Core/GameProject/CGameProject.h
@@ -14,6 +14,8 @@
 #include <Common/types.h>
 #include <Common/FileIO/CFileLock.h>
 
+namespace nod { class DiscWii; }
+
 class CGameProject
 {
     TString mProjectName;
@@ -56,9 +58,11 @@ class CGameProject
 public:
     ~CGameProject();
 
+
     bool Save();
     bool Serialize(IArchive& rArc);
     bool BuildISO(const TString& rkIsoPath, IProgressNotifier *pProgress);
+    bool MergeISO(const TString& rkIsoPath, nod::DiscWii *pOriginalIso, IProgressNotifier *pProgress);
     void GetWorldList(std::list<CAssetID>& rOut) const;
     CAssetID FindNamedResource(const TString& rkName) const;
     CPackage* FindPackage(const TString& rkName) const;
@@ -100,7 +104,7 @@ public:
     inline float BuildVersion() const                   { return mBuildVersion; }
     inline bool IsWiiBuild() const                      { return mBuildVersion >= 3.f; }
     inline bool IsTrilogy() const                       { return mGame <= eCorruption && mBuildVersion >= 3.593f; }
-    inline bool IsWiiDeAsobu() const                    { return mGame < eCorruption && mBuildVersion >= 3.f && mBuildVersion < 3.593f; }
+    inline bool IsWiiDeAsobu() const                    { return mGame <= eCorruption && mBuildVersion >= 3.f && mBuildVersion < 3.593f; }
 };
 
 #endif // CGAMEPROJECT_H
diff --git a/src/Editor/CProjectSettingsDialog.cpp b/src/Editor/CProjectSettingsDialog.cpp
index fbac8c96..1082de49 100644
--- a/src/Editor/CProjectSettingsDialog.cpp
+++ b/src/Editor/CProjectSettingsDialog.cpp
@@ -8,6 +8,9 @@
 #include <Common/AssertMacro.h>
 #include <Core/GameProject/CGameExporter.h>
 #include <Core/GameProject/COpeningBanner.h>
+
+#include <nod/nod.hpp>
+
 #include <QFileDialog>
 #include <QFuture>
 #include <QFutureWatcher>
@@ -147,12 +150,58 @@ void CProjectSettingsDialog::BuildISO()
 
     if (!IsoPath.isEmpty())
     {
+        bool NeedsDiscMerge = pProj->IsWiiDeAsobu() || pProj->IsTrilogy();
+        std::unique_ptr<nod::DiscBase> pBaseDisc = nullptr;
+
+        if (NeedsDiscMerge)
+        {
+            FilterString += ";*.wbfs";
+            QString SourceIsoPath = UICommon::OpenFileDialog(this, "Select the original ISO", FilterString, DefaultPath);
+
+            if (SourceIsoPath.isEmpty())
+                return;
+
+            // Verify this ISO matches the original
+            bool IsWii;
+            pBaseDisc = nod::OpenDiscFromImage(*TO_TWIDESTRING(SourceIsoPath), IsWii);
+
+            if (!pBaseDisc || !IsWii)
+            {
+                UICommon::ErrorMsg(this, "The ISO provided is not a valid Wii ISO!");
+                return;
+            }
+
+            const nod::Header& rkHeader = pBaseDisc->getHeader();
+            TString GameID = pProj->GameID();
+
+            if (strncmp(*GameID, rkHeader.m_gameID, 6) != 0)
+            {
+                UICommon::ErrorMsg(this, "The ISO provided doesn't match the project!");
+                return;
+            }
+        }
+
         if (gpEdApp->CookAllDirtyPackages())
         {
             CProgressDialog Dialog("Building ISO", false, true, this);
             Dialog.DisallowCanceling();
-            QFuture<void> Future = QtConcurrent::run(pProj, &CGameProject::BuildISO, TO_TSTRING(IsoPath), &Dialog);
-            Dialog.WaitForResults(Future);
+            bool Success;
+
+            if (!NeedsDiscMerge)
+            {
+                QFuture<bool> Future = QtConcurrent::run(pProj, &CGameProject::BuildISO, TO_TSTRING(IsoPath), &Dialog);
+                Success = Dialog.WaitForResults(Future);
+            }
+            else
+            {
+                QFuture<bool> Future = QtConcurrent::run(pProj, &CGameProject::MergeISO, TO_TSTRING(IsoPath), (nod::DiscWii*) pBaseDisc.get(), &Dialog);
+                Success = Dialog.WaitForResults(Future);
+            }
+
+            if (Success)
+                UICommon::InfoMsg(this, "Success", "ISO built successfully!");
+            else
+                UICommon::ErrorMsg(this, "ISO build failed!");
         }
     }
 }