Scripts/DK: correctly handle Blood Tap

src/server/scripts/Spells/spell_dk.cpp

@@ -31,10 +31,12 @@
 #include "PlayerAI.h"
 #include "Spell.h"
 #include "SpellAuraEffects.h"
+#include "SpellAuras.h"
 #include "SpellHistory.h"
 #include "SpellMgr.h"
 #include "SpellScript.h"
 #include "TemporarySummon.h"
+#include "Unit.h"
 
 enum DeathKnightSpells
 {
@@ -3049,6 +3051,126 @@ public:
     }
 };
 
+#define DKBloodTapScriptName "spell_dk_blood_tap"
+
+// 45529 - Blood Tap
+class spell_dk_blood_tap : public SpellScriptLoader
+{
+public:
+    spell_dk_blood_tap() : SpellScriptLoader(DKBloodTapScriptName) { }
+
+    class spell_dk_blood_tap_AuraScript : public AuraScript
+    {
+        PrepareAuraScript(spell_dk_blood_tap_AuraScript);
+
+    public:
+        spell_dk_blood_tap_AuraScript()
+        {
+            _runeIndex = MAX_RUNES;
+        }
+
+        void SetRuneIndex(uint8 index)
+        {
+            _runeIndex = index;
+        }
+
+        void HandleApply(AuraEffect const* aurEff, AuraEffectHandleModes /*mode*/)
+        {
+            PreventDefaultAction();
+
+            Player* player = GetTarget()->ToPlayer();
+            if (!player)
+                return;
+
+            if (player->GetClass() != CLASS_DEATH_KNIGHT || _runeIndex == MAX_RUNES)
+                return;
+
+            player->AddRuneByAuraEffect(_runeIndex, RUNE_DEATH, aurEff);
+        }
+
+        void Register() override
+        {
+            OnEffectApply += AuraEffectApplyFn(spell_dk_blood_tap_AuraScript::HandleApply, EFFECT_1, SPELL_AURA_CONVERT_RUNE, AURA_EFFECT_HANDLE_REAL);
+        }
+
+    private:
+        uint8 _runeIndex;
+    };
+
+    AuraScript* GetAuraScript() const override
+    {
+        return new spell_dk_blood_tap_AuraScript();
+    }
+
+    class spell_dk_blood_tap_SpellScript : public SpellScript
+    {
+        PrepareSpellScript(spell_dk_blood_tap_SpellScript);
+
+    public:
+        spell_dk_blood_tap_SpellScript()
+        {
+            _runeIndex = MAX_RUNES;
+        }
+
+        void HandleEffect(SpellEffIndex effIndex)
+        {
+            PreventHitDefaultEffect(effIndex);
+
+            Unit* caster = GetCaster();
+            if (caster->GetTypeId() != TYPEID_PLAYER)
+                return;
+
+            Player* player = caster->ToPlayer();
+            if (player->GetClass() != CLASS_DEATH_KNIGHT)
+                return;
+
+            // needed later
+            if (Spell* spell = GetSpell())
+                spell->SetRuneState(caster->ToPlayer()->GetRunesState());
+
+            uint8 resetIndex;
+            // Rune reset:
+            // If both runes are on cooldown, reset the shorter one
+            // If only one rune is on cooldown, reset that rune
+            if (!player->GetRuneCooldown(1))
+                resetIndex = 0; // 1 is ready, so reset 0 (no matter if it's on cd)
+            else if (!player->GetRuneCooldown(0) || player->GetRuneCooldown(1) < player->GetRuneCooldown(0))
+                resetIndex = 1; // 0 is ready, or both are on cd and 1 is shorter, so reset 1
+            else
+                resetIndex = 0; // both are on cd and 0 is shorter, reset 0
+
+            // if both runes are the same type, transform the same one as above
+            if (player->GetCurrentRune(0) == player->GetCurrentRune(1))
+                _runeIndex = resetIndex;
+            else // otherwise transform the blood rune
+                _runeIndex = player->GetCurrentRune(0) == RUNE_BLOOD ? 0 : 1;
+
+            player->SetRuneCooldown(resetIndex, 0);
+        }
+
+        void SetRuneIndex(SpellEffIndex /*effIndex*/)
+        {
+            if (Aura* aura = GetHitAura())
+                if (spell_dk_blood_tap_AuraScript* script = aura->GetScript<spell_dk_blood_tap_AuraScript>(DKBloodTapScriptName))
+                    script->SetRuneIndex(_runeIndex);
+        }
+
+        void Register() override
+        {
+            OnEffectLaunch += SpellEffectFn(spell_dk_blood_tap_SpellScript::HandleEffect, EFFECT_0, SPELL_EFFECT_ACTIVATE_RUNE);
+            OnEffectHitTarget += SpellEffectFn(spell_dk_blood_tap_SpellScript::SetRuneIndex, EFFECT_1, SPELL_EFFECT_APPLY_AURA);
+        }
+
+    private:
+        uint8 _runeIndex;
+    };
+
+    SpellScript* GetSpellScript() const override
+    {
+        return new spell_dk_blood_tap_SpellScript();
+    }
+};
+
 void AddSC_deathknight_spell_scripts()
 {
     new spell_dk_acclimation();
@@ -3104,4 +3226,5 @@ void AddSC_deathknight_spell_scripts()
     new spell_dk_raise_ally_initial();
     new spell_dk_raise_ally();
     new spell_dk_ghoul_thrash();
+    new spell_dk_blood_tap();
 }