OmicronFrames/src/types/unitframe.lua

-- Copyright 2023 <omicron.me@protonmail.com>
--
-- This file is part of Omicron Frames
-- 
-- Omicron Frames is free software: you can redistribute it and/or modify it
-- under the terms of the GNU General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option)
-- any later version.
--
-- Omicron Frames is distributed in the hope that it will be useful, but
-- WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
-- or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
-- more details.
--
-- You should have received a copy of the GNU General Public License along with
-- Omicron Frames. If not, see <https://www.gnu.org/licenses/>. 
local omif = select(2, ...)
local types = omif.GetModule("types")

---@class StatusBar
local StatusBar = types.StatusBar

---@class AuraList
local AuraList = types.AuraList

---@class SquareIndicator
local SquareIndicator = types.SquareIndicator

---@class AuraTrigger
local AuraTrigger = types.AuraTrigger

---@class StatusTrigger
local StatusTrigger = types.StatusTrigger

---@class UnitFrame
local UnitFrame = types.CreateClass("UnitFrame")
types.UnitFrame = UnitFrame


local colors = {
    hostile = {0.5, 0.0, 0.0},
    neutral = {0.7, 0.7, 0.0},
    healthy = {0, 0.1, 0},
    high = {0, 1.0, 0},
    mid = {1.0, 1.0, 0},
    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}
}

---@param unit string
---@param config table
function UnitFrame:Init(unit, config)
    local width = config.size.width
    local height = config.size.height

    self.unit = unit
    self.hideInRaid = config.hideInRaid

    self.rangeFriendly = config.range.friendly
    self.rangeEnemy = config.range.enemy
    self.rangeFade = config.range.fade

    local secure = self:CreateSecureFrame(width, height)
    secure:Hide()

    self:SetMouseBinds(config.mouse)

    self.hp = StatusBar:new(self, width, height, 0, true)
    self.power = StatusBar:new(self, width, 6, 2, false)
    self.power:Hide()
    self.auras = AuraList:new(self)

    local overlays = CreateFrame("Frame", nil, secure)
    overlays:SetFrameStrata("MEDIUM")
    overlays:SetFrameLevel(100)
    overlays:SetAllPoints(true)
    overlays:Show()
    self.overlays = overlays
    self:CreateName()

    self:CreateTriggers(config.triggers)

    self:RegisterEvents()
    self:Enable()
end

function UnitFrame:CreateTriggers(triggers)
    for _, trigger in ipairs(triggers) do
        local kind = trigger.kind
        local indicator = self:CreateIndicator(trigger.indicator)
        if kind == "AuraTrigger" then
            local at = AuraTrigger:new(
                    trigger.spellId,
                    trigger.own,
                    trigger.requiredCount,
                    trigger.invert
            )
            at:SetTarget(indicator)
            self.auras:AddTrigger(at)
        elseif kind == "StatusTrigger" then
            local st = StatusTrigger:new(trigger.status, trigger.requiredCount, trigger.invert)
            st:SetTarget(indicator)
            self.auras:AddTrigger(st)
        end
    end
end


function UnitFrame:CreateIndicator(indicator)
    local kind = indicator.kind
    if kind == "SquareIndicator" then
        return SquareIndicator:new(
            self,
            indicator.size,
            indicator.point,
            indicator.x,
            indicator.y,
            indicator.color
        )
    else
        error(string.format("Invalid Indicator kind `%s` requested", indicator.kind))
    end
end

function UnitFrame:StartRangeTicker()
    if self.rangeTicker then
        return
    end
    local delta = 0.45 + (fastrandom(0, 100)/1000)
    self.rangeTicker = C_Timer.NewTicker(delta, function()
        self:UpdateRange()
    end)
end

-- TODO: maybe string indicators need to be a thing
function UnitFrame:CreateName()
    local name = self.overlays:CreateFontString(nil, "BACKGROUND")
    self.name = name

    name:SetFont("Interface\\AddOns\\OmicronFrames\\media\\fonts\\roboto\\Roboto-Bold.ttf", 13, "")
    name:SetTextColor(0.8, 0.6, 0.1)
    name:SetShadowColor(0, 0, 0)
    name:SetShadowOffset(1, -1)
    name:SetPoint("CENTER")
    name:SetText("")
    name:Show()
end

function UnitFrame:StopRangeTicker()
    if self.rangeTicker then
        self.rangeTicker:Cancel()
        self.rangeTicker = nil
    end
end

-- Returns
--  true if the unit is in range
--  false if the unit is not in range
--  nil if it could not be determined 
function UnitFrame:IsInRange()
    local unit = self.unit
    local friendlySpell = self.rangeFriendly
    local enemySpell = self.rangeEnemy

    -- Prefer to use configured spells
    if friendlySpell and UnitCanAssist("player", unit) then
        return IsSpellInRange(friendlySpell, unit) == 1
    elseif enemySpell and UnitCanAttack("player", unit) then
        return IsSpellInRange(enemySpell, unit) == 1
    end
    
    -- Fall back to raid/party only range check 
    local inRange, checkedRange = UnitInRange(unit)
    if checkedRange then
        return inRange
    end

    return nil
end

function UnitFrame:UpdateRange()
    local unit = self.unit
    
    if UnitIsDeadOrGhost(unit) then
        self.secureFrame:SetAlpha(1.0)
        return
    end

    local inRange = self:IsInRange()
    if inRange or inRange == nil then
        self.secureFrame:SetAlpha(1.0)
    else
        self.secureFrame:SetAlpha(self.rangeFade)
    end
end

function UnitFrame:SetSpellAction(button, spell)
    local secure = self.secureFrame
    local attributeName = button:gsub("type", "spell")

    secure:SetAttribute(button, "spell")
    secure:SetAttribute(attributeName, spell)
end

function UnitFrame:SetMacroAction(button, macro)
    local secure = self.secureFrame
    local attributeName = button:gsub("type", "macrotext")

    secure:SetAttribute(button, "macro")
    secure:SetAttribute(attributeName, macro:gsub("@UNIT", "@" .. self.unit))
end

local function ModifiersToPrefix(mods)
    if not mods then
        return ""
    end

    local result = {}
    if mods.ctrl then
        table.insert(result, "CTRL-")
    end
    if mods.shift then
        table.insert(result, "SHIFT-")
    end
    if mods.alt then
        table.insert(result, "ALT-")
    end
    return table.concat(result, "")
end

function UnitFrame:PrepareWheelBinds(bindings)
    -- By default you can't use the SecureUnitButtonTemplate to handle scroll
    -- wheel "clicks". We solve this by using SecureHandler*Template to bind
    -- scroll wheel to the current unit frame button when the mouse enters and
    -- to unbind it when the mouse leaves. When the keybind is triggered the
    -- button will receive a click with a custom button name
    
    -- Build a table with all lines of the scroll wheel + modifier combinations
    -- we want to bind, then concat it with newlines into a full script
    local unit = self.unit
    local bindScript = {
        [[self:ClearBindings()]],
    }
    for _, bind in ipairs(bindings) do
        if bind.button == "wheel-up" or bind.button == "wheel-down" then
            local prefix = ModifiersToPrefix(bind.mods)
            local button = bind.button == "wheel-up" and "MOUSEWHEELUP" or "MOUSEWHEELDOWN"
            table.insert(bindScript,
                string.format([[self:SetBindingClick(false, "%s%s", "OmicronSecureFrame%s", "%s%s")]], 
                    prefix, button, unit, prefix, bind.button
                ) 
            )
        end
    end
    
    local secure = self.secureFrame
    local bindScript = table.concat(bindScript, "\n")

    local removeBindScript = [[
        self:ClearBindings()
    ]]
    secure:SetAttribute("_onenter", bindScript)
    secure:SetAttribute("_onleave", removeBindScript)
    secure:SetAttribute("_onshow", removeBindScript)
    secure:SetAttribute("_onhide", removeBindScript)
end

function UnitFrame:SetMouseBinds(binds)
    local secure = self.secureFrame
    local unit = self.unit

    secure:RegisterForClicks("AnyDown")
    self:PrepareWheelBinds(binds)
    
    for _, bind in ipairs(binds) do
        local prefix = ModifiersToPrefix(bind.mods)
        local button
        -- This is kind of a mess but there are two options for attribute names
        -- in the normal click case (mouse1, mouse2, etc) we have to put the
        -- modifiers before type so it would become SHIFT-type1 for mouse1 with
        -- shift
        --
        -- Since scroll wheel is not supported here and we manually create
        -- click events with global scroll wheel binds while we are moused over
        -- the unit frame they receive the modifiers after the type so it would become
        -- type-SHIFT-wheel-up for wheel-up with shift
        if bind.button:find("mouse", 1, true) == 1 then
            button = prefix .. bind.button:gsub("^mouse", "type")
        elseif bind.button:find("wheel", 1, true) == 1 then
            button = "type-" .. prefix .. bind.button
        else
            error("Keybinds were invalid")
        end
        if bind.kind == "macro" then
            self:SetMacroAction(button, bind.data)
        elseif bind.kind == "spell" then
            self:SetSpellAction(button, bind.data) 
        else
            secure:SetAttribute(button, bind.kind)
        end
    end
end

function UnitFrame:CreateSecureFrame(width, height)
    local name = "OmicronSecureFrame" .. self.unit
    local templates = table.concat({
        "SecureUnitButtonTemplate",
        "SecureHandlerEnterLeaveTemplate",
        "SecureHandlerShowHideTemplate",
    }, ",")
    local secure = CreateFrame("Button", name, UIParent, templates)
    self.secureFrame = secure
    secure:SetFrameStrata("MEDIUM")
    secure:SetFrameLevel(0)
    secure:SetAttribute("unit", self.unit)
    secure:SetSize(width, height)

    return secure
end

function UnitFrame:Enable()
    local secure = self.secureFrame
    if self.hideInRaid then
        local condition = "[@UNIT,exists,nogroup:raid] show; hide"
        condition = condition:gsub("@UNIT", "@" .. self.unit)
        RegisterAttributeDriver(secure, "state-visibility", condition)
    elseif self.unit ~= "player" then
        RegisterUnitWatch(self.secureFrame, false)
    else
        self.secureFrame:Show()
    end
    if UnitExists(self.unit) then
        self:UpdateAll(self:HasUnitChanged())
    end
end

function UnitFrame:Disable()
    if self.hideInRaid then
        UnregisterAttributeDriver(self.secureFrame, "state-visibility")
    elseif self.unit ~= "player" then
        UnregisterUnitWatch(self.secure_frame, self.condition)
    else
        self.secureFrame:Hide()
    end
end

-- returns true if this unit is a target unit
function UnitFrame:IsTargetFrame()
    return self.unit:match(".*target") == self.unit
end

-- returns the target owner
function UnitFrame:GetTargetOwner()
    if self.unit == "target" then
        return "player"
    else
        return self.unit:sub(1, -7)
    end
end

-- Set position from the center of the UIParent
function UnitFrame:SetPosition(x, y)
    self.secureFrame:SetPoint("CENTER", UIParent, x, y)
end

function UnitFrame:RegisterEvents()
    local unit = self.unit
    local secure = self.secureFrame

    secure:SetScript("OnShow", function(frame, event, ...)
        self:OnShow()
    end)

    secure:SetScript("OnHide", function(frame, event, ...)
        self:OnHide()
    end)

    secure:SetScript("OnEvent", function(frame, event, ...)
        self[event](self, ...)
    end)

    if self:IsTargetFrame() then
        local owner = self:GetTargetOwner()
        if owner == "player" then
            secure:RegisterEvent("PLAYER_TARGET_CHANGED")
        else
            secure:RegisterUnitEvent("UNIT_TARGET", self.parent:GetTargetOwner())
        end
    end

    secure:RegisterUnitEvent("UNIT_AURA", unit)
    secure:RegisterUnitEvent("UNIT_HEALTH", unit)
    secure:RegisterUnitEvent("UNIT_MAXHEALTH", unit)
    secure:RegisterUnitEvent("UNIT_NAME_UPDATE", unit)
end

-- returns whether or not the unit guid has changed since the last call to this
-- function. Always returns false if the current unit does not exist.
function UnitFrame:HasUnitChanged()
    local unit = self.unit
    if not UnitExists(unit) then
        self.guid = nil
        return false
    end

    local guid = UnitGUID(unit)
    if self.guid == guid then
        return false
    else
        self.guid = guid
        return true
    end
end

function UnitFrame:UNIT_AURA(unit, info)
    if self:HasUnitChanged() then
        self:UpdateAll(true)
        return
    end
    if self.auras:Update(info) then
        self:UpdateHealthColor()
    end
end

function UnitFrame:UNIT_HEALTH()
    if self:HasUnitChanged() then
        self:UpdateAll(true)
        return
    end
    self:UpdateHealth()
    self:UpdateHealthColor()
end

function UnitFrame:UNIT_MAXHEALTH()
    if self:HasUnitChanged() then
        self:UpdateAll(true)
        return
    end
    self:UpdateHealth()
    self:UpdateHealthColor()
end

function UnitFrame:UNIT_TARGET()
    self:UpdateAll(self:HasUnitChanged())
end

function UnitFrame:PLAYER_TARGET_CHANGED()
    self:UpdateAll(self:HasUnitChanged())
end

function UnitFrame:UNIT_NAME_UPDATE()
    self:UpdateName()
end

function UnitFrame:ROLE_CHANGED_INFORM(name, changer, old, new)
    if UnitName(self.unit) == name then
         self:UpdateRole(new)
    end
end

function UnitFrame:UNIT_POWER_UPDATE()
    self:UpdatePower()
end

function UnitFrame:UpdateRole(role)
    if role == nil then
        role = UnitGroupRolesAssigned(self.unit)
    end
    if role == "HEALER" then
        self:EnablePower(Enum.PowerType.Mana)
    elseif role == "TANK" and select(3, UnitClass(self.unit)) == 6 then
        self:EnablePower(Enum.PowerType.RunicPower)
    else
        self:DisablePower()
    end
end

function UnitFrame:UpdatePower()
    local power = UnitPower(self.unit, self.powerType)
    local max = UnitPowerMax(self.unit, self.powerType)
    self.power:SetRange(0, max)
    self.power:SetValue(power)
end

function UnitFrame:EnablePower(type)
    self.powerType = type
    local color = PowerBarColor[type]
    self.power:SetColor(color.r, color.g, color.b)
    self.power:Show()
    self.secureFrame:RegisterUnitEvent("UNIT_POWER_UPDATE", self.unit)
    self:UpdatePower()
end

function UnitFrame:DisablePower()
    self.powerType = nil
    self.secureFrame:UnregisterEvent("UNIT_POWER_UPDATE")
    self.power:Hide()
end

function UnitFrame:OnShow()
    self.secureFrame:RegisterEvent("ROLE_CHANGED_INFORM")
    self:StartRangeTicker()
    self:UpdateAll(self:HasUnitChanged())
end

function UnitFrame:OnHide()
    self.secureFrame:UnregisterEvent("ROLE_CHANGED_INFORM")
    self.guid = nil
    self:StopRangeTicker()
end

function UnitFrame:UpdateAll(unitChanged)
    if not UnitExists(self.unit) then
        return
    end
    if unitChanged then
        self.auras:Reset()
    end
    self:UpdateHealth()
    self:UpdateHealthColor()
    self:UpdateRange()
    self:UpdateName()
    self:UpdateRole() -- Also calls UpdatePower if power is visible
end

function UnitFrame:UpdateHealth()
    local unit = self.unit
    local val = UnitHealth(unit)
    local max = UnitHealthMax(unit)

    self.hp:SetRange(0, max)
    self.hp:SetValue(val)
end

function UnitFrame:UpdateHealthColor()
    local unit = self.unit
    local val = UnitHealth(unit)
    local max = UnitHealthMax(unit)

    if max == 0 then
        return
    end
    local pct = val / max
    local isFriend = UnitIsFriend("player", unit)

    if not isFriend then
        self.hp:SetColor(unpack(colors.hostile))
    elseif self.auras.statusCount["Immune"] then
        self.hp:SetColor(unpack(colors.immune))
    elseif self.auras.statusCount["Bomb"] then
        self.hp:SetColor(unpack(colors.bomb))
    elseif self.auras.statusCount["Magic"] then
        self.hp:SetColor(unpack(colors.magic))
    elseif self.auras.statusCount["Curse"] then
        self.hp:SetColor(unpack(colors.curse))
    elseif self.auras.statusCount["Poison"] then
        self.hp:SetColor(unpack(colors.poison))
    elseif self.auras.statusCount["Disease"] then
        self.hp:SetColor(unpack(colors.disease))
    elseif isFriend and UnitIsDead(unit) then
        self.hp:SetColor(unpack(colors.dead))
    elseif isFriend and pct >= 0.90 then
        self.hp:SetColor(unpack(colors.healthy))
    elseif isFriend and pct >= 0.75 then
        local progress = (pct - 0.75) / (0.90-0.75)
        self.hp:SetInterpolatedColor(colors.mid, colors.high, progress)
    else
        local progress = pct / 0.75
        -- Quadratic progress so we get to the dangerous colors faster as hp
        -- goes lower
        progress = progress * progress
        self.hp:SetInterpolatedColor(colors.low, colors.mid, progress)
    end
end

function UnitFrame:UpdateName()
    if UnitIsPlayer(self.unit) then
        local _, class = UnitClass(self.unit)
        -- TODO: UnitClass can return Unknown and the color can be nil. Having
        -- the unit exist but not fully loaded is possibly a state we want to
        -- handle more generally
        local color = RAID_CLASS_COLORS[class] or {r=1, g=1, b=1}
        self.name:SetTextColor(color.r, color.g, color.b)
    else
        self.name:SetTextColor(1, 1, 1)
    end
    self.name:SetText(UnitName(self.unit):sub(1, 5))
end

--[[
    -- UNIT_AURA
    -- UNIT_CLASSIFICATION_CHANGED
    -- UNIT_COMBAT
    -- UNIT_CONNECTION
    -- UNIT_DISPLAYPOWER
    -- UNIT_FACTION
    -- UNIT_FLAGS
    -- UNIT_LEVEL
    -- UNIT_MANA
    -- UNIT_HEALTH_PREDICTION
    -- UNIT_PHASE
]]--