Add a status bar for incoming heals

This shows a prediction of incoming direct heals.

src/frames.lua

@@ -26,6 +26,8 @@ local colors = {
     disease = { 0.4, 0.2, 0.0 },
     poison = { 0.0, 0.7, 0.7 },
     curse = { 0.7, 0.0, 0.7 },
+    immune = { 0.0, 0.2, 0.4 },
+    bomb = { 1.0, 0.7, 0.7 },
     red = { 1, 0, 0 },
     blue = { 0.4, 0.4, 1 },
     light_blue = { 0.7, 0.7, 1 },
@@ -833,6 +835,51 @@ local function TriggerConfig()
     return triggers
 end
 
+local function StatusColorClassConfig()
+    local _, class = UnitClass("player")
+    if class == "DRUID" then
+        return {
+            { status = "Curse", color = colors.curse },
+            { status = "Poison", color = colors.poison },
+        }
+    elseif class == "MAGE" then
+        return {
+            { status = "Curse", color = colors.curse },
+        }
+    elseif class == "PALADIN" then
+        return {
+            { status = "Magic", color = colors.magic },
+            { status = "Poison", color = colors.poison },
+            { status = "Disease", color = colors.disease },
+        }
+    elseif class == "PRIEST" then
+        return {
+            { status = "Magic", color = colors.magic },
+            { status = "Disease", color = colors.disease },
+        }
+    elseif class == "SHAMAN" then
+        return {
+            { status = "Poison", color = colors.poison },
+            { status = "Disease", color = colors.disease },
+        }
+    end
+
+    return {}
+end
+
+local function StatusColorConfig()
+    local statusColors = {
+        { status = "Immune", color = colors.immune },
+        { status = "Bomb", color = colors.bomb },
+    }
+
+    for _, entry in ipairs(StatusColorClassConfig()) do
+        table.insert(statusColors, entry)
+    end
+
+    return statusColors
+end
+
 local types = omif.GetModule("types")
 local UnitFrame = types.UnitFrame
 local UnitGroup = types.UnitGroup

src/types/statusbar.lua

@@ -90,3 +90,17 @@ function StatusBar:SetColor(r, g, b)
         self.bg:SetVertexColor(0.2 * r, 0.2 * g, 0.2 * b)
     end
 end
+
+function StatusBar:SetAlpha(alpha)
+    self.bar:SetAlpha(alpha)
+end
+
+function StatusBar:GetColor()
+    local r, g, b = self.bar:GetStatusBarColor()
+    return r, g, b
+end
+
+function StatusBar:AttachTo(statusBar, point, relativePoint)
+    self.bar:ClearAllPoints()
+    self.bar:SetPoint(point, statusBar.bar:GetStatusBarTexture(), relativePoint)
+end

src/types/unitframe.lua

@@ -50,12 +50,6 @@ local colors = {
     low = { 1.0, 0.0, 0 },
     offline = { 0.7, 0.7, 0.7 },
     dead = { 0.7, 0.7, 0.7 },
-    magic = { 0.4, 0.4, 1.0 },
-    disease = { 0.4, 0.2, 0.0 },
-    poison = { 0.0, 0.7, 0.7 },
-    curse = { 0.7, 0.0, 0.7 },
-    immune = { 0.0, 0.2, 0.4 },
-    bomb = { 1.0, 0.7, 0.7 },
     cyan = { 0.0, 0.8, 0.8 },
     white = { 1.0, 1.0, 1.0 },
     shield = { 0.85, 0.95, 1.0 },
@@ -90,6 +84,8 @@ function UnitFrame:Init(unit, config)
     self.rangeEnemy = config.range.enemy
     self.rangeFade = config.range.fade
 
+    self.statusColors = config.statusColors
+
     local secure = self:CreateSecureFrame(width, height)
     secure:Hide()
 
@@ -101,14 +97,18 @@ function UnitFrame:Init(unit, config)
 
     self.hp = StatusBar:new(self, width, height, 0, true, texture, texture)
 
-    self.power = StatusBar:new(self, width, 6, 2, false, texture, texture)
+    self.power = StatusBar:new(self, width, 6, 3, false, texture, texture)
     self.power:Hide()
 
-    self.absorb = StatusBar:new(self, width, height, 1, true, texture_top)
+    self.incomingHeals = StatusBar:new(self, width, height, 1, true, texture)
+    self.incomingHeals:AttachTo(self.hp, "TOPLEFT", "TOPRIGHT")
+    self.incomingHeals:Hide()
+
+    self.absorb = StatusBar:new(self, width, height, 2, true, texture_top)
     self.absorb:SetColor(unpack(colors.absorb))
     self.absorb:SetFillStyle(Enum.StatusBarFillStyle.Reverse)
 
-    self.shield = StatusBar:new(self, width, height, 1, true, texture_bot)
+    self.shield = StatusBar:new(self, width, height, 2, true, texture_bot)
     self.shield:SetColor(unpack(colors.shield))
     self.shield:SetRange(0, 1)
     self.shield:SetValue(0.6)
@@ -448,6 +448,7 @@ function UnitFrame:RegisterEvents()
     secure:RegisterUnitEvent("UNIT_NAME_UPDATE", unit)
     secure:RegisterUnitEvent("UNIT_ABSORB_AMOUNT_CHANGED", unit)
     secure:RegisterUnitEvent("UNIT_HEAL_ABSORB_AMOUNT_CHANGED", unit)
+    secure:RegisterUnitEvent("UNIT_HEAL_PREDICTION", unit)
 end
 
 -- returns whether or not the unit guid has changed since the last call to this
@@ -508,6 +509,7 @@ function UnitFrame:UNIT_HEALTH()
     end
     self:UpdateHealth()
     self:UpdateHealthColor()
+    self:UpdateIncomingHeals()
 end
 
 function UnitFrame:UNIT_MAXHEALTH()
@@ -517,6 +519,15 @@ function UnitFrame:UNIT_MAXHEALTH()
     end
     self:UpdateHealth()
     self:UpdateHealthColor()
+    self:UpdateIncomingHeals()
+end
+
+function UnitFrame:UNIT_HEAL_PREDICTION()
+    if self:HasUnitChanged() then
+        self:UpdateAll(true)
+        return
+    end
+    self:UpdateIncomingHeals()
 end
 
 function UnitFrame:UNIT_TARGET()
@@ -541,10 +552,20 @@ function UnitFrame:UNIT_POWER_UPDATE()
     self:UpdatePower()
 end
 
-function UnitFrame:UpdateRole(role)
+function UnitFrame:GetRole(role)
     if role == nil then
         role = UnitGroupRolesAssigned(self.unit)
     end
+
+    if GetPartyAssignment("MAINTANK", self.unit) then
+        role = "TANK"
+    end
+
+    return role
+end
+
+function UnitFrame:UpdateRole(role)
+    role = self:GetRole(role)
     if role == "HEALER" then
         self:EnablePower(Enum.PowerType.Mana)
     elseif role == "TANK" and select(3, UnitClass(self.unit)) == 6 then
@@ -597,6 +618,7 @@ function UnitFrame:UpdateAll(unitChanged)
     end
     self:UpdateHealth()
     self:UpdateHealthColor()
+    self:UpdateIncomingHeals()
     self:UpdateAbsorb()
     self:UpdateShield()
     self:UpdateRange()
@@ -652,23 +674,24 @@ function UnitFrame:UpdateHealthColor()
 
     if not isFriend then
         self.hp:SetColor(unpack(colors.hostile))
-    elseif self.auras.statusCount["Immune"] then
+        return
-        self.hp:SetColor(unpack(colors.immune))
+    end
-    elseif self.auras.statusCount["Bomb"] then
+
-        self.hp:SetColor(unpack(colors.bomb))
+    local statusColor = nil
-    elseif self.auras.statusCount["Magic"] then
+    for _, entry in ipairs(self.statusColors) do
-        self.hp:SetColor(unpack(colors.magic))
+        if self.auras.statusCount[entry.status] then
-    elseif self.auras.statusCount["Curse"] then
+            statusColor = entry.color
-        self.hp:SetColor(unpack(colors.curse))
+            break
-    elseif self.auras.statusCount["Poison"] then
+        end
-        self.hp:SetColor(unpack(colors.poison))
+    end
-    elseif self.auras.statusCount["Disease"] then
+
-        self.hp:SetColor(unpack(colors.disease))
+    if statusColor then
-    elseif isFriend and UnitIsDead(unit) then
+        self.hp:SetColor(unpack(statusColor))
+    elseif UnitIsDead(unit) then
         self.hp:SetColor(unpack(colors.dead))
-    elseif isFriend and pct >= 0.90 then
+    elseif pct >= 0.90 then
         self.hp:SetColor(unpack(colors.healthy))
-    elseif isFriend and pct >= 0.75 then
+    elseif pct >= 0.75 then
         local progress = (pct - 0.75) / (0.90 - 0.75)
         self.hp:SetInterpolatedColor(colors.mid, colors.high, progress)
     else
@@ -693,3 +716,21 @@ function UnitFrame:UpdateName()
     end
     self.name:SetText(utf8_truncate((UnitName(self.unit)), 5))
 end
+
+function UnitFrame:UpdateIncomingHeals()
+    local incoming = UnitGetIncomingHeals(self.unit) or 0
+    local health = UnitHealth(self.unit)
+    local maxHealth = UnitHealthMax(self.unit)
+    incoming = math.min(incoming, maxHealth - health)
+
+    if incoming <= 0 then
+        self.incomingHeals:Hide()
+    else
+        local r, g, b = self.hp:GetColor()
+        self.incomingHeals:SetColor(r, g, b)
+        self.incomingHeals:SetAlpha(0.4)
+        self.incomingHeals:SetRange(0, maxHealth)
+        self.incomingHeals:SetValue(incoming)
+        self.incomingHeals:Show()
+    end
+end

src/types/unitgroup.lua

@@ -53,8 +53,8 @@ local function UnitFrameDefaultCompare(lhs, rhs)
         end
 
         -- tank < healer < damage < ?
-        local lrole = RoleValue(UnitGroupRolesAssigned(lunit))
+        local lrole = RoleValue(lhs:GetRole())
-        local rrole = RoleValue(UnitGroupRolesAssigned(runit))
+        local rrole = RoleValue(rhs:GetRole())
 
         if lrole ~= rrole then
             return lrole < rrole