From ebd09f196cbfa4517b38a1073022afbee84d4619 Mon Sep 17 00:00:00 2001
From: Paul Chote <pchote@users.noreply.github.com>
Date: Sun, 29 Mar 2015 18:50:20 +0100
Subject: [PATCH] Add upgrade support to SupportPower.

---
 .../Traits/SupportPowers/SupportPower.cs      |  7 +++----
 .../SupportPowers/SupportPowerManager.cs      | 20 +++++++++++++++----
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/OpenRA.Mods.Common/Traits/SupportPowers/SupportPower.cs b/OpenRA.Mods.Common/Traits/SupportPowers/SupportPower.cs
index 6edb273fdc..a1fb516512 100644
--- a/OpenRA.Mods.Common/Traits/SupportPowers/SupportPower.cs
+++ b/OpenRA.Mods.Common/Traits/SupportPowers/SupportPower.cs
@@ -12,7 +12,7 @@ using OpenRA.Traits;
 
 namespace OpenRA.Mods.Common.Traits
 {
-	public abstract class SupportPowerInfo : ITraitInfo
+	public abstract class SupportPowerInfo : UpgradableTraitInfo, ITraitInfo
 	{
 		[Desc("Measured in seconds.")]
 		public readonly int ChargeTime = 0;
@@ -53,15 +53,14 @@ namespace OpenRA.Mods.Common.Traits
 		public SupportPowerInfo() { OrderName = GetType().Name + "Order"; }
 	}
 
-	public class SupportPower
+	public class SupportPower : UpgradableTrait<SupportPowerInfo>
 	{
 		public readonly Actor Self;
-		public readonly SupportPowerInfo Info;
 		protected RadarPing ping;
 
 		public SupportPower(Actor self, SupportPowerInfo info)
+			: base(info)
 		{
-			Info = info;
 			Self = self;
 		}
 
diff --git a/OpenRA.Mods.Common/Traits/SupportPowers/SupportPowerManager.cs b/OpenRA.Mods.Common/Traits/SupportPowers/SupportPowerManager.cs
index 395833905e..6f49b1c926 100644
--- a/OpenRA.Mods.Common/Traits/SupportPowers/SupportPowerManager.cs
+++ b/OpenRA.Mods.Common/Traits/SupportPowers/SupportPowerManager.cs
@@ -132,7 +132,7 @@ namespace OpenRA.Mods.Common.Traits
 			if (!Powers.TryGetValue(key, out sp))
 				return;
 
-			sp.Disabled = false;
+			sp.PrerequisitesAvailable(true);
 		}
 
 		public void PrerequisitesUnavailable(string key)
@@ -141,7 +141,7 @@ namespace OpenRA.Mods.Common.Traits
 			if (!Powers.TryGetValue(key, out sp))
 				return;
 
-			sp.Disabled = true;
+			sp.PrerequisitesAvailable(false);
 			sp.RemainingTime = sp.TotalTime;
 		}
 
@@ -158,17 +158,25 @@ namespace OpenRA.Mods.Common.Traits
 		public int RemainingTime;
 		public int TotalTime;
 		public bool Active { get; private set; }
-		public bool Disabled { get; set; }
+		public bool Disabled { get { return !prereqsAvailable || !upgradeAvailable; } }
 
 		public SupportPowerInfo Info { get { return Instances.Select(i => i.Info).FirstOrDefault(); } }
 		public bool Ready { get { return Active && RemainingTime == 0; } }
 
+		bool upgradeAvailable;
+		bool prereqsAvailable = true;
+
 		public SupportPowerInstance(string key, SupportPowerManager manager)
 		{
 			this.manager = manager;
 			this.key = key;
 		}
 
+		public void PrerequisitesAvailable(bool available)
+		{
+			prereqsAvailable = available;
+		}
+
 		static bool InstanceDisabled(SupportPower sp)
 		{
 			return sp.Self.TraitsImplementing<IDisable>().Any(d => d.Disabled);
@@ -178,6 +186,10 @@ namespace OpenRA.Mods.Common.Traits
 		bool notifiedReady;
 		public void Tick()
 		{
+			upgradeAvailable = Instances.Any(i => !i.IsTraitDisabled);
+			if (!upgradeAvailable)
+				RemainingTime = TotalTime;
+
 			Active = !Disabled && Instances.Any(i => !i.Self.IsDisabled());
 			if (!Active)
 				return;
@@ -226,7 +238,7 @@ namespace OpenRA.Mods.Common.Traits
 			notifiedCharging = notifiedReady = false;
 
 			if (Info.OneShot)
-				Disabled = true;
+				PrerequisitesAvailable(false);
 		}
 	}