[WotLK][TC][Eluna][AIO] Custom class powerType / resource

--[[INSTRUCTIONS:
    1   Requires core modification to fix stance/form changing:
        in Unit.cpp
        void Unit::setPowerType(Powers new_powertype)
        change:
            case POWER_HAPPINESS:
                SetMaxPower(POWER_HAPPINESS, uint32(std::ceil(GetCreatePowers(POWER_HAPPINESS) * powerMultiplier)));
                SetPower(POWER_HAPPINESS, uint32(std::ceil(GetCreatePowers(POWER_HAPPINESS) * powerMultiplier)));
                break;
        to:
            case POWER_HAPPINESS:
                if (GetTypeId() == TYPEID_PLAYER && GetMaxPower(POWER_HAPPINESS) >= 0)
                    break;
                SetMaxPower(POWER_HAPPINESS, uint32(std::ceil(GetCreatePowers(POWER_HAPPINESS) * powerMultiplier)));
                SetPower(POWER_HAPPINESS, uint32(std::ceil(GetCreatePowers(POWER_HAPPINESS) * powerMultiplier)));
                break;

        You may want to apply other core modifications to give POWER_HAPPINESS as neutral of a behavior as possible and script
        functionality specific to each custom powerType when a specific player event occurs, or add some class checks in the
        core when a player's POWER_HAPPINESS is manipulated.

    2   If you want a class to use a custom resource, set the SERVER SIDE powertype for that class to 4 (HAPPINESS)
        in ChrClasses.dbc (the field that has a default value of 6 for Death Knights). Client side changes to ChrClasses.dbc
        are not required. Happiness is simply the powertype we use behind the scenes for all custom powertypes.

    3   Next, add an entry for your custom powertype in the customPowerTypes table. Each entry should have a table of classes,
        an id, a powerString, a bar color table, a defaultMax value and a consumingSpells table (for correct GetSpellInfo powerType).
        defaultMax should be the value you want, multiplied by 1000. All Happiness values have to be multiplied by 1000.

    4   To create ANY spell (in Spell.dbc) that uses or generates your custom powertype, give it a power cost or energize effect
        for powerType 4 (HAPPINESS). Make sure Happiness values (costs, energize effects etc) are multiplied by 1000.


TROUBLESHOOTING:
        Make sure that the class has the powerType set to 4 (HAPPINESS) in SERVER SIDE ChrClasses.dbc.
        Make sure that the class is listed in only one customPowerTypes[POWERTOKEN].classes table.
        Make sure table entries are separated by commas.
        MAKE SURE ALL RESOURCE VALUES ARE MULTIPLIED BY 1000. Happiness for some reason needs everything
        multiplied by 1000. That includes spell costs and generation or if you're setting current power value
        or maxvalue. All of them need to be multiplied by 1000.
        Custom spells should always use 4 (HAPPINESS) in Spell.dbc, but everywhere in-game the powerType and
        powerToken will match your custom values. Don't get confused and give a spell in Spell.dbc a custom
        cost powerType (eg. a Spell.dbc entry that costs powerType 201). It should instead have cost powerType 4 (HAPPINESS).
        Check the worldserver console for errors (scripts get loaded at the end).
        If you get UI errors in-game, keep in mind the line number might not be accurate if you have server code and
        client code in the same script file. Try subtracting the server-only line count, or split the script
        into a server file and client file.


Currently this system has some (minor?) flaws:
        - No error speech or text if you fail client-side resource check (not even a "Not enough happiness" message).
        - Spells linked in chat simply show the cost as "X Happiness", but it should be possible to alter the tooltip if the
          spell is found in a custom powertype's consumingSpells table.
        - The official Blizzard UI function UnitFrameManaBar_UpdateType(manaBar) from UnitFrame.lua is overwritten, which
          I guess is bad form because if the function is overwritten by two different addons, one of them will misbehave.
        - A class can't swap between two custom primary resources on the fly (like a druid swapping between mana, energy, rage)
          because the custom resource is using Happiness and MaxHappiness to store values. You would need to store extra values
          somewhere and change how unitframes set their manabar and change how spells' usable status is determined.
]]--


local AIO = AIO or require("AIO")

local customPowerTypes = {}                     -- Table for custom class resource data, used by server and client. Not all of it is used by both the server and client, but it's convenient to have everything in one place.


customPowerTypes["BRAINCELLS"] =                -- Our first custom class resource, used by Shamans and Rogues (remember to set ChrClasses.dbc powerTypes to use this example)
{
    classes = { "Shaman", "Rogue" },            -- List of classes that use this as primary resource. Remember to set powertype to 4 (HAPPINESS) in server's ChrClasses.dbc
    defaultMax = 300000,                        -- Happiness values need to be multiplied by 1000, so for the power bar to read "powertype 0 / 123", set defaultMax to 123000
    id = 200,                                   -- The id for our custom resource (first value returned by UnitPowerType). MUST BE UNIQUE. Some values below 200 are reserved for existing powerTypes.
    powerString = "Braincells",                 -- The string to be displayed for our powertype
    color = { r = 0.83, g = 0.76, b = 0.60 },   -- RGB colors for our power bar. Values should be between 0 and 1
    consumingSpells = { 90031 }                 -- Spell IDs that consume this resource. Required for correct GetSpellInfo(spellID) API functionality
}

customPowerTypes["TEMPER"] =                    -- Our second custom class resource, used by Warriors
{
    classes = { "Warrior" },
    defaultMax = 100000,
    id = 201,
    powerString = "Temper",
    color = { r = 1.0, g = 1.0, b = 1.0 },
    consumingSpells = { 90032 },
    passiveUpdate = {                                                                   -- Optional passive update data, containing update interval and an update function. Event gets automatically registered, so long as you use these names ("interval" and "UpdateFunc")
                        interval = 100,                                                 -- Update interval in milliseconds
                        UpdateFunc = function(event, delay, repeats, player)            -- power update function that gets called every [interval] milliseconds
                            -- points (re/de)generation per second
                            local regen = -10000
                            if (player:IsInCombat()) then
                                regen = 10000
                            end
                            local prevPower = player:GetPower(4)
                            local newPower = prevPower + (regen * (delay / 1000))
                            local maxPower = player:GetMaxPower(4)

                            if (newPower < 0) then
                                newPower = 0
                            elseif ( newPower > maxPower ) then
                                newPower = maxPower
                            end

                            player:SetPower(4, newPower)
                        end
                    }
}


if AIO.AddAddon() then
    local function OnLogin(event, player)
        local playerclass = player:GetClassAsString()

        for powerToken, powerTable in pairs(customPowerTypes) do
            for _, class in pairs(powerTable.classes) do
                if (playerclass == class) then
                    player:SetMaxPower(4, powerTable.defaultMax)

                    -- Prevent tiny mana bar from re-appearing below the main resource bar by setting max mana to 0 (normally appears for druids while they're shapeshifted)
                    -- Ideally you would set this class to never get any mana or maxmana increases in class levelup data
                    player:SetMaxPower(0, 0)

                    if (powerTable.passiveUpdate ~= nil) then
                        player:RegisterEvent(powerTable.passiveUpdate.UpdateFunc, powerTable.passiveUpdate.interval, 0)
                    end
                end
            end
        end

    end

    -- Prevent player's maximum Happiness from being reset to core default 1050 (aka 1050000). Preferably you would remove the reset in the core, because
    -- this way you can't change a player's maximum value without resetting it to the defaultMax, unless you store and restore desired maxpower per character
    local function OnLevelChange(event, player, oldLevel)
        local playerclass = player:GetClassAsString()
        for _, powerTable in pairs(customPowerTypes) do
            for _, class in pairs(powerTable.classes) do
                if (playerclass == class) then
                    if (powerTable.defaultMax ~= nil) then
                        player:SetMaxPower(4, powerTable.defaultMax)
                    end
                end
            end
        end
    end

    RegisterPlayerEvent( 3, OnLogin )
    RegisterPlayerEvent( 13, OnLevelChange )
    return
end


HAPPINESS_COST = "%d Happiness";
local myclass = UnitClass("player")

for customToken, powerTable in pairs(customPowerTypes) do
    _G[customToken] = powerTable.powerString
    PowerBarColor[customToken] = powerTable.color
    PowerBarColor[powerTable.id] = PowerBarColor[customToken]
end


-- Hook official Blizzard UI function GetSpellInfo(spellId or spellName or spellLink), to fix custom powerType return, if applicable
-- SpellLink as argument is broken in the official API, apparently.
local origGetSpellInfo = GetSpellInfo;
GetSpellInfo = function(...)
    local spellID = ...;
    local name, rank, icon, cost, isFunnel, powerType, castTime, minRange, maxRange = origGetSpellInfo(...)

    -- if a spellName was passed as argument, jump through some hoops to get the spellID
    if (type(spellID) == "string") then
        spellLink = GetSpellLink(spellID)
        if (spellLink ~= nil) then
            spellID = tonumber(strmatch(spellLink, "^\124c%x+\124Hspell:(%d+)\124h%[.*%]"))
        end
    end

    if (spellID ~= nil) then
        for _, powerTable in pairs(customPowerTypes) do
            local earlyReturn = false
            for _, id in pairs(powerTable.consumingSpells) do
                if (id == spellID) then
                    powerType = powerTable.id
                    break
                end
            end
            if earlyReturn then break end
        end
    end

    return name, rank, icon, cost, isFunnel, powerType, castTime, minRange, maxRange;
end

-- Hook official Blizzard UI function UnitPower(unit, powerType), to add custom value return, if applicable
local origUnitPower = UnitPower;
UnitPower = function(...)
    local unit, powerType = ...;

    if (UnitIsPlayer(unit)) then
        --local unitClass = UnitClass(unit)
        for customToken, powerTable in pairs(customPowerTypes) do
            if (powerTable.id == powerType) then
                -- ask for Happiness value instead
                return origUnitPower(unit, 4)
            end
        end
    end

    return origUnitPower(...);
end

-- Hook official Blizzard UI function UnitPowerMax(unit, powerType), to add custom value return, if applicable
local origUnitPowerMax = UnitPowerMax;
UnitPowerMax = function(...)
    local unit, powerType = ...;

    if (UnitIsPlayer(unit)) then
        for customToken, powerTable in pairs(customPowerTypes) do
            if (powerTable.id == powerType) then
                -- ask for Happiness value instead
                return origUnitPowerMax(unit, 4)
            end
        end
    end

    return origUnitPowerMax(...);
end

-- Hook official Blizzard UI function UnitPowerType(unit), to add custom powerType and powerToken return, if applicable
local origUnitPowerType = UnitPowerType;
UnitPowerType = function(...)
    local unit = ...;

    local unitClass = UnitClass(unit)
    for customToken, powerTable in pairs(customPowerTypes) do
        for _, class in pairs(powerTable.classes) do
            if (unitClass == class) then
                local powerType, _, altR, altG, altB = origUnitPowerType(unit)
                powerType = powerTable.id
                return powerType, customToken, altR, altG, altB
            end
        end
    end

    return origUnitPowerType(...);
end

-- Overwrite official Blizzard UI function UnitFrameManaBar_UpdateType(manaBar), hacky fix to display Happiness value and maxvalue in mana bar
function UnitFrameManaBar_UpdateType(manaBar)
    if ( not manaBar ) then
        return;
    end
    local unitFrame = manaBar:GetParent();
    local powerType, powerToken, altR, altG, altB = UnitPowerType(manaBar.unit);
    local prefix = _G[powerToken];
    local info = PowerBarColor[powerToken];
    if ( info ) then
        if ( not manaBar.lockColor ) then
            manaBar:SetStatusBarColor(info.r, info.g, info.b);
        end
    else
        if ( not altR) then
            -- couldn't find a power token entry...default to indexing by power type or just mana if we don't have that either
            info = PowerBarColor[powerType] or PowerBarColor["MANA"];
        else
            if ( not manaBar.lockColor ) then
                manaBar:SetStatusBarColor(altR, altG, altB);
            end
        end
    end

    -- We still want to use Happiness behind the scenes, so if we're using a custom resource, restore the powerType ID to 4: HAPPINESS
    if (customPowerTypes[powerToken] ~= nil) then
        powerType = 4;
    end
    manaBar.powerType = powerType;

    -- Update the manabar text
    if ( not unitFrame.noTextPrefix ) then
        SetTextStatusBarTextPrefix(manaBar, prefix);
    end
    TextStatusBar_UpdateTextString(manaBar);

    -- Setup newbie tooltip
    -- FIXME: Fix this to use powerToken instead of powerType
    if ( manaBar.unit ~= "pet" or powerToken == "HAPPINESS" ) then
        if ( unitFrame:GetName() == "PlayerFrame" ) then
            manaBar.tooltipTitle = prefix;
            manaBar.tooltipText = _G["NEWBIE_TOOLTIP_MANABAR_"..powerType];
        else
            manaBar.tooltipTitle = nil;
            manaBar.tooltipText = nil;
        end
    end
end

local function CustomPrimaryResourceTooltip_OnShow(self, ...)
    local _, _, spellID = GameTooltip:GetSpell()
    if (spellID == nil) then
        return
    end

    local numLines = GameTooltip:NumLines();
    local i = 1;

    for currentLine=1, numLines do
        local line = {};
        local left = _G["GameTooltipTextLeft"..currentLine];
        if ( left ) then
            line.w = true;
            line.leftR, line.leftG, line.leftB = left:GetTextColor();
            local t = left:GetText();

            local happinessCost = strmatch(t, "(%d+) Happiness")
            if (happinessCost ~= nil) then
                local _, powerToken = UnitPowerType("player")
                left:SetText(happinessCost.." ".._G[powerToken])
            end

            line.left = t;
        end
        i = i + 1;
    end

    GameTooltip:Show();
end


for powerToken, powerTable in pairs(customPowerTypes) do
    for _, class in pairs(powerTable.classes) do
        if (myclass == class) then
            UnitFrameManaBar_UpdateType(PlayerFrameManaBar)
            -- Hide mini mana bar
            AlternatePowerBar_UpdatePowerType(PlayerFrameAlternateManaBar)
        end
    end
end

GameTooltip:HookScript( "OnShow", CustomPrimaryResourceTooltip_OnShow )