From 285bafdbf46016dfe5a87720d4d4d0ae42f7c565 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tomasz=20Zieli=C5=84ski?=
 <tomaszzielinskijedenzwielu@gmail.com>
Date: Mon, 15 Jan 2024 07:41:29 +0100
Subject: [PATCH] Do not ban heroes in Prisons, they might be rehired after
 they are defeated. Restore hero to pool if Prison fails to be placed.

---
 lib/rmg/CMapGenerator.cpp                 | 11 -----------
 lib/rmg/CMapGenerator.h                   |  4 +---
 lib/rmg/modificators/PrisonHeroPlacer.cpp |  8 +++-----
 lib/rmg/modificators/PrisonHeroPlacer.h   |  2 +-
 lib/rmg/modificators/TreasurePlacer.cpp   |  3 ++-
 mapeditor/mapcontroller.cpp               |  5 +++++
 6 files changed, 12 insertions(+), 21 deletions(-)

diff --git a/lib/rmg/CMapGenerator.cpp b/lib/rmg/CMapGenerator.cpp
index 3af14b544..a41bc4b78 100644
--- a/lib/rmg/CMapGenerator.cpp
+++ b/lib/rmg/CMapGenerator.cpp
@@ -509,17 +509,6 @@ void CMapGenerator::banQuestArt(const ArtifactID & id)
 	map->getMap(this).allowedArtifact.erase(id);
 }
 
-void CMapGenerator::banHero(const HeroTypeID & id)
-{
-	map->getMap(this).banHero(id);
-}
-
-void CMapGenerator::unbanHero(const HeroTypeID & id)
-{
-	map->getMap(this).unbanHero(id);
-}
-
-
 Zone * CMapGenerator::getZoneWater() const
 {
 	for(auto & z : map->getZones())
diff --git a/lib/rmg/CMapGenerator.h b/lib/rmg/CMapGenerator.h
index fe5b5dde0..6561c8bc6 100644
--- a/lib/rmg/CMapGenerator.h
+++ b/lib/rmg/CMapGenerator.h
@@ -65,9 +65,7 @@ public:
 	const std::vector<ArtifactID> & getAllPossibleQuestArtifacts() const;
 	const std::vector<HeroTypeID> getAllPossibleHeroes() const;
 	void banQuestArt(const ArtifactID & id);
-	void banHero(const HeroTypeID& id);
-	void unbanHero(const HeroTypeID & id);
-
+	
 	Zone * getZoneWater() const;
 	void addWaterTreasuresInfo();
 
diff --git a/lib/rmg/modificators/PrisonHeroPlacer.cpp b/lib/rmg/modificators/PrisonHeroPlacer.cpp
index 06f02aba2..d4787784d 100644
--- a/lib/rmg/modificators/PrisonHeroPlacer.cpp
+++ b/lib/rmg/modificators/PrisonHeroPlacer.cpp
@@ -56,20 +56,18 @@ HeroTypeID PrisonHeroPlacer::drawRandomHero()
 		RandomGeneratorUtil::randomShuffle(allowedHeroes, zone.getRand());
         HeroTypeID ret = allowedHeroes.back();
         allowedHeroes.pop_back();
-
-        generator.banHero(ret);
 		return ret;
 	}
 	else
 	{
-		throw rmgException("No quest heroes left for prisons!");
+		throw rmgException("No unused heroes left for prisons!");
 	}
 }
 
-void PrisonHeroPlacer::unbanHero(const HeroTypeID & hid)
+void PrisonHeroPlacer::restoreDrawnHero(const HeroTypeID & hid)
 {
 	RecursiveLock lock(externalAccessMutex);
-	generator.unbanHero(hid);
+	allowedHeroes.push_back(hid);
 }
 
 VCMI_LIB_NAMESPACE_END
diff --git a/lib/rmg/modificators/PrisonHeroPlacer.h b/lib/rmg/modificators/PrisonHeroPlacer.h
index fa25c7b33..4f5cc13c6 100644
--- a/lib/rmg/modificators/PrisonHeroPlacer.h
+++ b/lib/rmg/modificators/PrisonHeroPlacer.h
@@ -27,7 +27,7 @@ public:
 
 	int getPrisonsRemaning() const;
 	HeroTypeID drawRandomHero();
-	void unbanHero(const HeroTypeID & hid);
+	void restoreDrawnHero(const HeroTypeID & hid);
 
 private:
     void getAllowedHeroes();
diff --git a/lib/rmg/modificators/TreasurePlacer.cpp b/lib/rmg/modificators/TreasurePlacer.cpp
index 5ea775573..cb63d927d 100644
--- a/lib/rmg/modificators/TreasurePlacer.cpp
+++ b/lib/rmg/modificators/TreasurePlacer.cpp
@@ -130,7 +130,8 @@ void TreasurePlacer::addAllPossibleObjects()
 				HeroTypeID hid = prisonHeroPlacer->drawRandomHero();
 				oi.destroyObject = [hid, prisonHeroPlacer]()
 				{
-					prisonHeroPlacer->unbanHero(hid);
+					// Hero can be used again
+					prisonHeroPlacer->restoreDrawnHero(hid);
 				};
 
 				auto factory = VLC->objtypeh->getHandlerFor(Obj::PRISON, 0);
diff --git a/mapeditor/mapcontroller.cpp b/mapeditor/mapcontroller.cpp
index 00dcaad1c..659f01b05 100644
--- a/mapeditor/mapcontroller.cpp
+++ b/mapeditor/mapcontroller.cpp
@@ -131,7 +131,12 @@ void MapController::repairMap(CMap * map) const
 		//fix hero instance
 		if(auto * nih = dynamic_cast<CGHeroInstance*>(obj.get()))
 		{
+			// All heroes present on map or in prisons need to be allowed to rehire them after they are defeated
+
+			// FIXME: How about custom scenarios where defeated hero cannot be hired again?
+
 			map->allowedHeroes.insert(nih->getHeroType());
+
 			auto type = VLC->heroh->objects[nih->subID];
 			assert(type->heroClass);
 			//TODO: find a way to get proper type name