coa-omen/Omen.lua

-----------------------------------------------------------------------------
-- Addon declaration
local Omen = LibStub("AceAddon-3.0"):NewAddon("Omen", "AceConsole-3.0", "AceEvent-3.0", "AceTimer-3.0")
local L = LibStub("AceLocale-3.0"):GetLocale("Omen", false)
_G["Omen"] = Omen

-- TODO: Add TPS calculations
-- TODO: Add config
-- TODO: Add visible/audible warnings
-- TODO: Add Sink/SharedMedia stuff
-- TODO: Code optimizations


-----------------------------------------------------------------------------
-- Localize some global functions
local floor, format = floor, format
local sort = sort
local UnitDetailedThreatSituation = UnitDetailedThreatSituation
local UnitExists, UnitGUID, UnitName, UnitClass = UnitExists, UnitGUID, UnitName, UnitClass
local UnitIsPlayer, UnitPlayerControlled, UnitCanAttack = UnitIsPlayer, UnitPlayerControlled, UnitCanAttack
local GetNumRaidMembers, GetNumPartyMembers = GetNumRaidMembers, GetNumPartyMembers


-----------------------------------------------------------------------------
-- Local variables used
local db
local defaults = {
	profile = {
		Scale        = 100,
		GrowUp       = false,
		Autocollapse = false,
		NumBars      = 10,
		Locked       = false,
		BarHeight    = 16,
		ShowWith = {
			Pet = true,
			Alone = false,
			Party = true,
			Raid = true,
			Resting = false,
			PVP = false,
			Dungeon = true
		},
		Classes = {
			["DEATHKNIGHT"] = true,
			["DRUID"] = true,
			["HUNTER"] = true,
			["MAGE"] = true,
			["PALADIN"] = true,
			["PET"] = true,
			["PRIEST"] = true,
			["ROGUE"] = true,
			["SHAMAN"] = true,
			["WARLOCK"] = true,
			["WARRIOR"] = true
		},
	},
}
local guidNameLookup = {}   -- Format: guidNameLookup[guid] = "Unit Name"
local guidClassLookup = {}  -- Format: guidClassLookup[guid] = "CLASS"
local timers = {}           -- Format: timers.timerName = timer returned from AceTimer-3.0
local bars = {}             -- Format: bars[i] = frame containing the i-th bar from the top of Omen
local inRaid, inParty       -- boolean variables indicating if the player is in a raid and/or party

Omen.GuidNameLookup = guidNameLookup
Omen.GuidClassLookup = guidClassLookup
Omen.Timers = timers
Omen.Bars = bars

local default_color = {     -- Default bar color for units not in the player's raid/party
	r = 1,
	g = 0,
	b = 0,
	a = 0.6
}
local pet_color = {         -- Default pet color for all pets in the player's raid/party
	r = 0.83,               -- This is a light green-ish color, or sort of a washed out
	g = 0.98,               -- version of the Hunter color
	b = 0.74,
	a = 0.6
}

-- For speedups. Rather than concantenating every time we need a unitID, we just look
-- it up instead. That is rID[i] is much faster than format("raid%d", i) or "raid"..i
local pID = {}
local ptID = {}
local ppID = {}
local pptID = {}
local rID = {}
local rtID = {}
local rpID = {}
local rptID = {}
for i = 1, 4 do
	pID[i] = format("party%d", i)
	ptID[i] = format("party%dtarget", i)
	ppID[i] = format("partypet%d", i)
	pptID[i] = format("partypet%dtarget", i)
end
for i = 1, 40 do
	rID[i] = format("raid%d", i)
	rtID[i] = format("raid%dtarget", i)
	rpID[i] = format("raidpet%d", i)
	rptID[i] = format("raidpet%dtarget", i)
end


-----------------------------------------------------------------------------
-- Omen initialization and frame functions

local function startmoving()
	if not db.Locked then
		Omen.Anchor:StartMoving()
	end
end
local function stopmoving()
	Omen.Anchor:StopMovingOrSizing()
	Omen:SetAnchors()
	Omen:UpdateBars()
end
local function sizing()
	Omen:ResizeBars()
	Omen:UpdateBars()
end

function Omen:OnInitialize()
	-- Create savedvariables
	self.db = LibStub("AceDB-3.0"):New("Omen3DB", defaults)
	self.db.RegisterCallback(self, "OnProfileChanged", "OnProfileChanged")
	self.db.RegisterCallback(self, "OnProfileCopied", "OnProfileChanged")
	self.db.RegisterCallback(self, "OnProfileReset", "OnProfileChanged")
	db = self.db.profile

	-- Create anchor
	self.Anchor = CreateFrame("Frame", "OmenAnchor", UIParent)
	self.Anchor:SetResizable(true)
	self.Anchor:SetMinResize(90, db.BarHeight)
	self.Anchor:SetMovable(true)
	self.Anchor:SetPoint("CENTER", UIParent, "CENTER")
	self.Anchor:SetWidth(225)
	self.Anchor:SetHeight(150)
	self.Anchor:SetClampedToScreen(true)
	self.Anchor:SetScript("OnHide", stopmoving)

	-- Create Title
	self.Title = CreateFrame("Frame", "OmenTitle", self.Anchor)
	self.Title:SetPoint("TOPLEFT", self.Anchor, "TOPLEFT")
	self.Title:SetPoint("TOPRIGHT", self.Anchor, "TOPRIGHT")
	self.Title:SetHeight(16)
	self.Title:EnableMouse(true)
	self.Title:SetScript("OnMouseDown", startmoving)
	self.Title:SetScript("OnMouseUp", stopmoving)
	self.Title.Texture = self.Title:CreateTexture()
	self.Title.Texture:SetAllPoints()
	self.Title.Texture:SetTexture(0, 0, 0, 0.8)

	-- Create Title text
	self.TitleText = self.Title:CreateFontString(nil, nil, "GameFontNormal")
	self.TitleText:SetPoint("LEFT", self.Title, "LEFT", 10, 0)
	self.TitleText:SetJustifyH("LEFT")
	self.TitleText:SetTextColor(1, 1, 1, 0.95)
	self.defaultTitle = "Omen|cffffcc003|r"
	self.TitleText:SetText(self.defaultTitle)

	-- Create Bar List
	self.BarList = CreateFrame("Frame", "OmenBarList", self.Anchor)
	self.BarList:SetResizable(true)
	self.BarList:EnableMouse(true)
	self.BarList:SetPoint("TOPLEFT", self.Title, "BOTTOMLEFT")
	self.BarList:SetPoint("TOPRIGHT", self.Title, "BOTTOMRIGHT")
	self.BarList:SetPoint("BOTTOMLEFT", self.Anchor, "BOTTOMLEFT")
	self.BarList:SetPoint("BOTTOMRIGHT", self.Anchor, "BOTTOMRIGHT")
	self.BarList.Texture = self.BarList:CreateTexture()
	self.BarList.Texture:SetAllPoints()
	self.BarList.Texture:SetTexture(0, 0, 0, 0.8)
	self.BarList:SetScript("OnMouseDown", startmoving)
	self.BarList:SetScript("OnMouseUp", stopmoving)

	-- Create resizing corner grip
	self.Grip = CreateFrame("Button", "OmenResizeGrip", self.BarList)
	self.Grip:SetNormalTexture("Interface\\AddOns\\Omen\\ResizeGrip")
	self.Grip:SetHighlightTexture("Interface\\AddOns\\Omen\\ResizeGrip")
	self.Grip:SetWidth(16)
	self.Grip:SetHeight(16)
	self.Grip:SetPoint("BOTTOMRIGHT", self.BarList, "BOTTOMRIGHT", 0, 1)
	if db.Locked then
		self.Grip:Hide()
	end
	self.Grip:SetScript("OnMouseDown", function()
		if not db.Locked then
			self.Anchor.IsMovingOrSizing = true
			self.Anchor:SetScript("OnSizeChanged", sizing)
			self.Anchor:StartSizing()
		end
	end)
	self.Grip:SetScript("OnMouseUp", function()
		self.Anchor:SetScript("OnSizeChanged", nil)
		self.Anchor:StopMovingOrSizing()
		self.Anchor.IsMovingOrSizing = nil
		sizing()
		self:SetAnchors()
	end)
	self.Grip:SetScript("OnHide", self.Grip:GetScript("OnHide"))
	
	-- Register events that are "always on". The rest are in :OnEnable()
	self:RegisterEvent("PLAYER_LOGIN")
	self:RegisterEvent("PLAYER_ENTERING_WORLD")
end

function Omen:PLAYER_LOGIN()
	-- We set up anchors here because we only want to do it once on
	-- PLAYER_LOGIN, hence we don't do it in OnEnable() which triggers on
	-- the same event as well as on every subsequent Enable()/Disable() calls.
	-- It cannot be earlier than PLAYER_LOGIN because layout-cache.txt
	-- is loaded just before this event fires.
	self:SetAnchors(true)
	self:UnregisterEvent("PLAYER_LOGIN")

	--[[
	-- Optional launcher support for LDB-1.1 if present, this code is placed here so
	-- that it runs after all other addons have loaded since we don't embed LDB-1.1
	local DataBroker = LibStub("LibDataBroker-1.1", true)
	if DataBroker then
		local launcher = DataBroker:NewDataObject("Omen", {
		    type = "launcher",
		    icon = "Interface\\AddOns\\Omen\\icon",
		    OnClick = function(clickedframe, button)
				if button == "RightButton" then Omen:ShowConfig() else Omen:Toggle() end
			end,
		})
	end
	]]
end

function Omen:OnEnable()
	self:RegisterEvent("UNIT_THREAT_LIST_UPDATE")
	self:RegisterEvent("UNIT_THREAT_SITUATION_UPDATE")
	self:RegisterEvent("PLAYER_TARGET_CHANGED")
	self:RegisterEvent("UNIT_PET")
	self:RegisterEvent("PARTY_MEMBERS_CHANGED")

	self:PARTY_MEMBERS_CHANGED()
	self:PLAYER_TARGET_CHANGED()
	self:Toggle(true)
end

function Omen:Disable()
	self:UnregisterEvent("UNIT_THREAT_LIST_UPDATE")
	self:UnregisterEvent("UNIT_THREAT_SITUATION_UPDATE")
	self:UnregisterEvent("PLAYER_TARGET_CHANGED")
	self:UnregisterEvent("UNIT_PET")
	self:UnregisterEvent("PARTY_MEMBERS_CHANGED")

	-- Cancel all timers (well at least nil them all
	-- out in timers[], since AceTimer-3.0 cancels
	-- them all OnDisable anyway).
	for k, v in pairs(timers) do
		self:CancelTimer(v, true)
		timers[k] = nil
	end

	self:Toggle(false)
end

function Omen:OnProfileChanged(db, name)
	db = self.db.profile
end

function Omen:SetAnchors(useDB)
	local x, y, w, h

	-- Set the scale, since the scaling affects the position
	self.Anchor:SetScale(db.Scale / 100)

	-- Get position
	if useDB then
		x, y = db.PositionX, db.PositionY
		if not x or not y then
			self.Anchor:ClearAllPoints()
			self.Anchor:SetPoint("CENTER", UIParent, "CENTER", 0, 0)
			x, y = self.Anchor:GetLeft(), db.GrowUp and self.Anchor:GetBottom() or self.Anchor:GetTop()
		end
	else
		x, y = self.Anchor:GetLeft(), db.GrowUp and self.Anchor:GetBottom() or self.Anchor:GetTop()
	end

	-- Get width/height
	w = useDB and db.PositionW or self.Anchor:GetWidth()
	h = useDB and db.PositionH or self.Anchor:GetHeight()

	-- Set the anchors and size
	self.Anchor:ClearAllPoints()
	self.Anchor:SetPoint(db.GrowUp and "BOTTOMLEFT" or "TOPLEFT", UIParent, "BOTTOMLEFT", x, y)
	self.Anchor:SetWidth(w)
	self.Anchor:SetHeight(h)
	self.Anchor:SetUserPlaced(nil)

	-- Save the data
	db.PositionX, db.PositionY = x, y
	db.PositionW, db.PositionH = w, h
end

function Omen:Toggle(setting)
	if setting == nil then
		setting = not self.Anchor:IsShown()
	end
	if setting then
		self.Anchor:Show()
	else
		self.Anchor:Hide()
	end
end


-----------------------------------------------------------------------------
-- Omen bar stuff

setmetatable(bars, {__index = function(self, barID)
	local bar = CreateFrame("Frame", nil, Omen.BarList)
	self[barID] = bar

	bar:SetWidth(Omen.BarList:GetWidth())
	bar:SetHeight(db.BarHeight)
	bar:SetPoint("TOPLEFT", Omen.BarList, "TOPLEFT", 0, (barID-1) * -db.BarHeight)

	bar.Text1 = bar:CreateFontString(nil, nil, "GameFontNormalSmall")
	bar.Text1:SetPoint("LEFT", bar, "LEFT", 10, 0)
	bar.Text1:SetJustifyH("LEFT")
	bar.Text1:SetTextColor(1, 1, 1, 1)
	bar.Text1:SetText("abc")
	bar.Text2 = bar:CreateFontString(nil, nil, "GameFontNormalSmall")
	bar.Text2:SetPoint("RIGHT", bar, "RIGHT", -10, 0)
	bar.Text2:SetJustifyH("RIGHT")
	bar.Text2:SetTextColor(1, 1, 1, 1)
	bar.Text2:SetText("def")

	bar.texture = bar:CreateTexture()
	bar.texture:SetTexture(1, 1, 1, 0.6)
	bar.texture:SetPoint("TOPLEFT", bar, "TOPLEFT")
	bar.texture:SetPoint("BOTTOMLEFT", bar, "BOTTOMLEFT")
	return bar
end})


-----------------------------------------------------------------------------
-- Omen event functions

-- Fired when a mob has its threat list updated. The mob that
-- had its list updated is the first parameter of the event.
function Omen:UNIT_THREAT_LIST_UPDATE(event, unitID)
	if unitID == "target" then
		self:UpdateBars()
	end
end

-- Fired when a unit's threat situation changes. The unit that
-- had a change in threat situation is the first parameter of
-- the event. Note that this only triggers when major state
-- changes, not when the raw threat values change.
function Omen:UNIT_THREAT_SITUATION_UPDATE(...)
	self:UpdateBars()
end

function Omen:PLAYER_TARGET_CHANGED()
	-- Stop our unit update timer for updating threat on "targettarget"
	if timers.UpdateBars then
		self:CancelTimer(timers.UpdateBars, true)
		timers.UpdateBars = nil
	end

	if UnitExists("target") then
		self:UpdateBars()
	else
		self:ClearAll()
	end
end

function Omen:UNIT_PET()
	self:PARTY_MEMBERS_CHANGED()
	-- TODO: Add automatic show/hide triggers
end

function Omen:PLAYER_ENTERING_WORLD()
	if UnitGUID("player") then
		guidNameLookup[UnitGUID("player")] = UnitName("player")
	end
	-- TODO: Add automatic show/hide triggers
end

local lastPartyUpdateTime = GetTime()

function Omen:PARTY_MEMBERS_CHANGED()
	inRaid = GetNumRaidMembers() > 0
	inParty = GetNumPartyMembers() > 0

	-- Run the update if the last call is more than 0.5 seconds ago else
	-- schedule an update 0.5 seconds later if one isn't already scheduled
	if GetTime() - lastPartyUpdateTime > 0.5 then
		self:UpdatePartyGUIDs()
	elseif not timers.UpdatePartyGUIDs then
		timers.UpdatePartyGUIDs = self:ScheduleTimer("UpdatePartyGUIDs", 0.5)
	end
end

-- This function updates the name and class guid lookup tables of the raid
function Omen:UpdatePartyGUIDs()
	lastPartyUpdateTime = GetTime()
	if timers.UpdatePartyGUIDs then
		self:CancelTimer(timers.UpdatePartyGUIDs, true)
		timers.UpdatePartyGUIDs = nil
	end

	local _
	for k, v in pairs(guidClassLookup) do
		guidClassLookup[k] = nil
	end
	_, guidClassLookup[UnitGUID("player")] = UnitClass("player")
	if UnitExists("pet") then
		guidClassLookup[UnitGUID("pet")] = "PET"
		guidNameLookup[UnitGUID("pet")] = UnitName("pet")--.." ["..UnitName("player").."]"
	end

	if inParty or inRaid then
		local playerFmt = inRaid and rID or pID
		local petFmt = inRaid and rpID or ppID
		local currentPartySize = inRaid and GetNumRaidMembers() or GetNumPartyMembers()
		
		for i = 1, currentPartySize do
			local unitID = playerFmt[i]
			local pGUID = UnitGUID(unitID)
			
			if pGUID then
				guidNameLookup[pGUID] = UnitName(unitID)
				_, guidClassLookup[pGUID] = UnitClass(unitID)
				
				-- lookup pet (if existing)
				local petID = petFmt[i]
				local petGUID = UnitGUID(petID)
				if petGUID then
					guidNameLookup[petGUID] = UnitName(petID)--.." ["..UnitName(unitID).."]"
					guidClassLookup[petGUID] = "PET"
				end
			end
		end
	end
end


-----------------------------------------------------------------------------
-- Omen update functions

--[[
First, some definitions:
* mob - enemy creature
* threat list - a mob's list of possible targets, along with each possible target's current threat value
* threat situation - the situation that a unit is currently in (either globally, or with respect to a certain mob)
* scaled percentage - a threat percentage, where 100% means you will pull aggro (become the primary target of the mob), and thus this % cannot be higher than 100% under normal circumstances
* raw threat percentage - the percentage of the units threat when divided by the threat of the mob's current primary target, this % CAN be over 100%
---------
state = UnitThreatSituation(unit, mob)

Returns the unit's threat situation with respect to the given mob. The state can be one of the following values:
nil = the unit is not on the mob's threat list
0 = 0-99% raw threat percentage (no indicator shown)
1 = 100% or more raw threat percentage (yellow warning indicator shown)
2 = tanking, other has 100% or more raw threat percentage (orange indicator shown)
3 = tanking, all others have less than 100% raw percentage threat (red indicator shown)
---------
state = UnitThreatSituation(unit)

Returns the unit's maximum threat state on any mob's threat list.
---------
isTanking, state, scaledPercent, rawPercent, threatValue = UnitDetailedThreatSituation(unit, mob)

Returns detailed information about the unit's state on the mob's threat list.
isTanking is true if the unit the primary target of the mob (and by definition has 100% threat)
state is the unit's threat situation, as listed above.
scaledPercent is the current percent threat of the unit, scaled in the 0-100% range based on distance from target.
rawPercent is the current percent threat of the unit relative to the primary target of the mob.
threatValue is the amount of threat that the unit has on the mob's threat list. This is roughly approximate to the amount of damage and healing the unit has done.
---------
r, g, b = GetThreatStatusColor(state)

Returns the colors used in the UI to represent each major threat state. 
]]

local threatTable = {} -- Format: threatTable[guid] = threatValue
local sortTable = {}   -- Format: threatTable[i] = guid -- used for sorting by sortfunction()
local tankGUID         -- Used to store which unit is tanking and hence has 100% threat by definition

local function sortfunction(a, b)
	return threatTable[a] > threatTable[b]
end

local function updatethreat(unitid, mobunitid)
	local guid = UnitGUID(unitid)
	if guid and not threatTable[guid] then
		local isTanking, state, scaledPercent, rawPercent, threatValue = UnitDetailedThreatSituation(unitid, mobunitid)
		threatTable[guid] = threatValue or -1
		if isTanking then tankGUID = guid end
	end
end

function Omen:UpdateBars()
	if not self.Anchor:IsShown() then return end
	
	-- TODO: Put a update throtle on this function

	-- Figure out which mob to show threat on, either "target" or "targettarget"
	-- It has to be attackable and not human controlled, otherwise return
	local mob = "target"
	if UnitIsPlayer(mob) or UnitPlayerControlled(mob) or not UnitCanAttack("player", mob) then
		mob = "targettarget"
	end
	if not UnitExists(mob) or UnitIsPlayer(mob) or UnitPlayerControlled(mob) or not UnitCanAttack("player", mob) then
		self:ClearAll()
		return
	end
	local mobGUID = UnitGUID(mob)
	guidNameLookup[mobGUID] = UnitName(mob)
	self.TitleText:SetText(guidNameLookup[mobGUID])

	-- Schedule a repeating timer for updating threat on "targettarget"
	-- since we get no events on a targettarget change.
	if mob == "targettarget" and not timers.UpdateBars then
		timers.UpdateBars = self:ScheduleRepeatingTimer("UpdateBars", 0.5)
	end

	-- We want the mob's target just in case the tank isn't
	-- in our raid (say an NPC or some other player)
	local mobTarget = mob.."target"
	local mobTargetGUID = UnitGUID(mobTarget)
	if mobTargetGUID and not guidNameLookup[mobTargetGUID] then
		guidNameLookup[mobTargetGUID] = UnitName(mobTarget)
	end

	-- Clear the threat table
	for k, v in pairs(threatTable) do
		threatTable[k] = nil
	end
	threatTable[mobGUID] = -1
	tankGUID = nil

	-- Get data for threat on mob by scanning the whole raid
	if inParty or inRaid then
		if inRaid then
			for i = 1, GetNumRaidMembers() do
				updatethreat(rID[i], mob)
				updatethreat(rpID[i], mob)
				updatethreat(rtID[i], mob)
				updatethreat(rptID[i], mob)
			end
		else
			for i = 1, GetNumPartyMembers() do
				updatethreat(pID[i], mob)
				updatethreat(ppID[i], mob)
				updatethreat(ptID[i], mob)
				updatethreat(pptID[i], mob)
			end
		end
		
	end
	if not inRaid then
		updatethreat("player", mob)
		updatethreat("pet", mob)
		updatethreat("target", mob)
		updatethreat("pettarget", mob)
	end
	updatethreat("focus", mob)
	updatethreat("focustarget", mob)
	updatethreat(mobTarget, mob)

	-- Sort the threatTable
	local i = 1
	for k, v in pairs(threatTable) do
		if v >= 0 then
			sortTable[i] = k
			i = i + 1
		end
	end
	for j = i, #sortTable do
		sortTable[j] = nil
	end
	if #sortTable == 0 then
		self:ClearAll()
		return
	end
	sort(sortTable, sortfunction)

	-- Now update the bars on screen
	local w = self.BarList:GetWidth()
	local h = self.BarList:GetHeight()
	local topthreat = threatTable[sortTable[1]]
	local tankThreat = threatTable[tankGUID or mobTargetGUID or sortTable[1]]
	i = 1
	for j = 1, #sortTable do
		local guid = sortTable[j]
		local class = guidClassLookup[guid]
		if class == nil or db.Classes[class] then -- class == nil implies its not in your raid
			local bar = bars[j]
			bar.Text1:SetText(guidNameLookup[guid] or L["Unknown"])
			bar.Text2:SetText(floor(threatTable[guid] / 100).." ["..floor(threatTable[guid] / tankThreat * 100).."%]")
			local c = (class == "PET" and pet_color) or RAID_CLASS_COLORS[class] or default_color
			bar.texture:SetVertexColor(c.r, c.g, c.b)
			bar.texture:SetWidth(w * threatTable[guid] / topthreat)
			bar:Show()
			i = i + 1
			if db.Autocollapse then
				if i > db.NumBars then break end
			else
				if floor(h / bar:GetHeight()) < i then break end
			end
		end
	end
	-- And hide the rest
	for j = i, #bars do
		bars[j]:Hide()
	end
	if db.Autocollapse then
		self.Anchor:SetHeight(i*db.BarHeight)
	end
end

function Omen:ClearAll()
	for i = 1, #bars do
		bars[i]:Hide()
	end
	self.TitleText:SetText(self.defaultTitle)
	if db.Autocollapse then
		self.Anchor:SetHeight(db.BarHeight)
	end
end

function Omen:ResizeBars()
	local w = Omen.BarList:GetWidth()
	for i = 1, #bars do
		bars[i]:SetWidth(w)
		bars[i]:SetHeight(db.BarHeight)
	end
end