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")

local NameFont = CreateFont("OicronFrames_NameFont")
NameFont:SetFont("Interface\\AddOns\\OmicronFrames\\media\\fonts\\roboto\\Roboto-Bold.ttf", 13, "")
NameFont:SetShadowColor(0, 0, 0)
NameFont:SetShadowOffset(1, -1)

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

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

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

---@class IconIndicator
local IconIndicator = types.IconIndicator

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

---@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 },
    cyan = { 0.0, 0.8, 0.8 },
    white = { 1.0, 1.0, 1.0 },
    shield = { 0.85, 0.95, 1.0 },
    absorb = { 0.5, 0.4, 0.75 },
}

-- This trucates _codepoints_ not graphemes. If combination codepoints are
-- contained in the string, it will not properly truncate and may return
-- incorrect graphemes
local function utf8_truncate(s, n)
    local init = 0
    for _ = 1, n do
        local _, last = string.find(s, ".[\128-\191]*", init + 1)
        if not last then
            return s
        end
        init = last
    end
    return string.sub(s, 1, init)
end

---@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

    self.statusColors = config.statusColors

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

    self:SetMouseBinds(config.mouse)

    local texture = "Interface\\Addons\\OmicronFrames\\media\\textures\\bar_subtle_diagonal"
    local texture_top = "Interface\\Addons\\OmicronFrames\\media\\textures\\bar_subtle_diagonal_top"
    local texture_bot = "Interface\\Addons\\OmicronFrames\\media\\textures\\bar_subtle_diagonal_bottom"

    self.hp = StatusBar:new(self, width, height, 0, true, texture, texture)

    self.power = StatusBar:new(self, width, 6, 2, false, texture, texture)
    self.power:Hide()

    self.absorb = StatusBar:new(self, width, height, 1, 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:SetColor(unpack(colors.shield))
    self.shield:SetRange(0, 1)
    self.shield:SetValue(0.6)
    self.shield:SetFillStyle(Enum.StatusBarFillStyle.Reverse)

    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 _, config in ipairs(triggers) do
        local indicator = self:CreateIndicator(config.indicator)
        ---@type Trigger
        local triggerType = types[config.kind]
        local trigger = triggerType.CreateFromConfig(self, config)
        trigger:SetTarget(indicator)
    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,
            indicator.showStacks,
            indicator.fadeTime,
            indicator.flashTime
        )
    elseif kind == "IconIndicator" then
        return IconIndicator:new(
            self,
            indicator.size,
            indicator.point,
            indicator.x,
            indicator.y,
            indicator.showStacks,
            indicator.fadeTime,
            indicator.flashTime
        )
    elseif kind == "BorderIndicator" then
        return BorderIndicator:new(self, indicator.thickness, 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:SetFontObject(NameFont)
    name:SetTextColor(0.8, 0.6, 0.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 C_Spell.IsSpellInRange(friendlySpell, unit) == true
    elseif enemySpell and UnitCanAttack("player", unit) then
        return C_Spell.IsSpellInRange(enemySpell, unit) == true
    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.secureFrame, false)
    end
    self.secureFrame:Hide()
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:RegisterEvent("GROUP_ROSTER_UPDATE")
    secure:RegisterUnitEvent("UNIT_AURA", unit)
    secure:RegisterUnitEvent("UNIT_HEALTH", unit)
    secure:RegisterUnitEvent("UNIT_MAXHEALTH", unit)
    secure:RegisterUnitEvent("UNIT_NAME_UPDATE", unit)
    secure:RegisterUnitEvent("UNIT_ABSORB_AMOUNT_CHANGED", unit)
    secure:RegisterUnitEvent("UNIT_HEAL_ABSORB_AMOUNT_CHANGED", 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_HEAL_ABSORB_AMOUNT_CHANGED()
    if self:HasUnitChanged() then
        self:UpdateAll(true)
        return
    end
    self:UpdateAbsorb()
end

function UnitFrame:UNIT_ABSORB_AMOUNT_CHANGED()
    if self:HasUnitChanged() then
        self:UpdateAll(true)
        return
    end
    self:UpdateShield()
end

function UnitFrame:GROUP_ROSTER_UPDATE()
    if self:HasUnitChanged() then
        self:UpdateAll(true)
        return
    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:UpdateAbsorb()
    self:UpdateShield()
    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:UpdateShield()
    local unit = self.unit
    local val = UnitGetTotalAbsorbs(unit)
    local max = UnitHealthMax(unit)
    if val <= 0 then
        self.shield:Hide()
    else
        self.shield:SetRange(0, max)
        self.shield:SetValue(val)
        self.shield:Show()
    end
end

function UnitFrame:UpdateAbsorb()
    local unit = self.unit
    local val = UnitGetTotalHealAbsorbs(unit)
    local max = UnitHealthMax(unit)
    if val <= 0 then
        self.absorb:Hide()
    else
        self.absorb:SetRange(0, max)
        self.absorb:SetValue(val)
        self.absorb:Show()
    end
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))
        return
    end

    local statusColor = nil
    for _, entry in ipairs(self.statusColors) do
        if self.auras.statusCount[entry.status] then
            statusColor = entry.color
            break
        end
    end

    if statusColor then
        self.hp:SetColor(unpack(statusColor))
    elseif UnitIsDead(unit) then
        self.hp:SetColor(unpack(colors.dead))
    elseif pct >= 0.90 then
        self.hp:SetColor(unpack(colors.healthy))
    elseif 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(utf8_truncate((UnitName(self.unit)), 5))
end