local Tinkr, Bastion = ...

local ElementalShamanModule = Bastion.Module:New('ElementalShaman')
local Player = Bastion.UnitManager:Get('player')
local Target = Bastion.UnitManager:Get('target')
local Pet = Bastion.UnitManager:Get('pet')


-- Initialize SpellBook
local SpellBook = Bastion.SpellBook:New()

-- Spells
local LightningBolt = SpellBook:GetSpell(188196)
local LavaBurst = SpellBook:GetSpell(51505)
local EarthShock = SpellBook:GetSpell(8042)
local FlameShock = SpellBook:GetSpell(188389)
local Stormkeeper = SpellBook:GetSpell(191634)
local ElementalBlast = SpellBook:GetSpell(117014)
local ChainLightning = SpellBook:GetSpell(188443)
local Earthquake = SpellBook:GetSpell(61882)
local FireElemental = SpellBook:GetSpell(198067)
local StormElemental = SpellBook:GetSpell(192249)
local LiquidMagmaTotem = SpellBook:GetSpell(192222)
local Icefury = SpellBook:GetSpell(210714)
local FrostShock = SpellBook:GetSpell(196840)
local Ascendance = SpellBook:GetSpell(114050)
local PrimordialWave = SpellBook:GetSpell(375982)
local AstralShift = SpellBook:GetSpell(108271)
local WindShear = SpellBook:GetSpell(57994)
local Skyfury = SpellBook:GetSpell(462854)

-- Buffs and Debuffs
local MasterOfTheElements = SpellBook:GetSpell(260734)
local SurgeOfPower = SpellBook:GetSpell(285514)
local LavaSurge = SpellBook:GetSpell(77762)
local Icefury_Buff = SpellBook:GetSpell(210714)
local FlameShock_Debuff = SpellBook:GetSpell(188389)
local MagmaChamber_Buff = SpellBook:GetSpell(381933)
local SplinteredElements_Buff = SpellBook:GetSpell(382043)
local WindGust_Buff = SpellBook:GetSpell(263806)
local EarthShield_Buff = SpellBook:GetSpell(383648)
local LightningShield_Buff = SpellBook:GetSpell(192106)
local PrimordialWave_Buff = SpellBook:GetSpell(375986)

-- Custom Units
local BestTarget = Bastion.UnitManager:CreateCustomUnit('besttarget', function()
    local bestTarget = nil
    local highestHP = 0

    Bastion.UnitManager:EnumEnemies(function(unit)
        if unit:IsAffectingCombat() and Player:IsWithinCombatDistance(unit, 40) and Player:CanSee(unit) and Player:IsFacing(unit) then
            local hp = unit:GetHP()
            if hp > highestHP then
                bestTarget = unit
                highestHP = hp
            end
        end
    end)

    return bestTarget or Target
end)

-- APLs
local DefaultAPL = Bastion.APL:New('default')
local AoEAPL = Bastion.APL:New('aoe')
local CooldownAPL = Bastion.APL:New('cooldown')
local DefensiveAPL = Bastion.APL:New('defensive')
local OutOfCombatAPL = Bastion.APL:New('outofcombat')

-- Helper Functions
local function ShouldAoE()
    local _, enemies = Bastion.UnitManager:GetEnemiesWithMostEnemies(10)
    return #enemies >= 2
end

-- Default APL
DefaultAPL:AddSpell(
    PrimordialWave:CastableIf(function(self)
        return BestTarget:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
    end):SetTarget(BestTarget)
)

DefaultAPL:AddSpell(
    LavaBurst:CastableIf(function(self)
        return BestTarget:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and (Player:GetAuras():FindAny(PrimordialWave_Buff):IsUp() or (Player:GetAuras():FindAny(LavaSurge):IsUp() and Player:IsMoving()))
    end):SetTarget(BestTarget)
)

DefaultAPL:AddSpell(
    EarthShock:CastableIf(function(self)
        return BestTarget:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and Player:GetPower() >= 60 and (Player:GetAuras():FindAny(MasterOfTheElements):IsUp() or (Player:GetPower() + self:GetCost()) > Player:GetMaxPower())
    end):SetTarget(BestTarget)
)

DefaultAPL:AddSpell(
    ElementalBlast:CastableIf(function(self)
        return BestTarget:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and Player:GetAuras():FindAny(MasterOfTheElements):IsUp()
    end):SetTarget(BestTarget)
)

DefaultAPL:AddSpell(
    FlameShock:CastableIf(function(self)
        return BestTarget:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and not BestTarget:GetAuras():FindAny(FlameShock_Debuff):IsUp()
    end):SetTarget(BestTarget)
)

DefaultAPL:AddSpell(
    LightningBolt:CastableIf(function(self)
        return BestTarget:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
    end):SetTarget(BestTarget)
)

DefaultAPL:AddSpell(
    FrostShock:CastableIf(function(self)
        return BestTarget:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling() and Player:IsMoving()
    end):SetTarget(BestTarget)
)

-- AoE APL
AoEAPL:AddSpell(
    ChainLightning:CastableIf(function(self)
        return BestTarget:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and ShouldAoE()
    end):SetTarget(BestTarget)
)

AoEAPL:AddSpell(
    Earthquake:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and Player:GetPower() >= 60 and ShouldAoE()
    end):SetTarget(Bastion.UnitManager:Get('none')):OnCast(function(self)
        local loc = Bastion.UnitManager:FindEnemiesCentroid(10, 40)
        if loc then
            self:Click(loc)
        end
    end)
)

-- Cooldown APL
CooldownAPL:AddSpell(
    StormElemental:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling() and not Pet:Exists()
    end):SetTarget(Player)
)

CooldownAPL:AddSpell(
    FireElemental:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling() and not StormElemental:IsKnownAndUsable()
    end):SetTarget(Player)
)

CooldownAPL:AddSpell(
    Stormkeeper:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
    end):SetTarget(Player)
)

CooldownAPL:AddSpell(
    LiquidMagmaTotem:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling() and ShouldAoE()
    end):SetTarget(Bastion.UnitManager:Get('none')):OnCast(function(self)
        local loc = Bastion.UnitManager:FindEnemiesCentroid(10, 40)
        if loc then
            self:Click(loc)
        end
    end)
)

CooldownAPL:AddSpell(
    Ascendance:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
    end):SetTarget(Player)
)

-- Defensive APL
DefensiveAPL:AddSpell(
    AstralShift:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and Player:GetHP() < 50
    end):SetTarget(Player)
)

-- Interrupt APL
local InterruptAPL = Bastion.APL:New('interrupt')
InterruptAPL:AddSpell(
    WindShear:CastableIf(function(self)
        return BestTarget:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and BestTarget:IsInterruptible()
    end):SetTarget(BestTarget)
)

-- Out of Combat APL

-- Module Sync
ElementalShamanModule:Sync(function()
    if Player:IsMounted() then
        return
    end

    DefensiveAPL:Execute()
    InterruptAPL:Execute()

    if Player:IsAffectingCombat() then
        CooldownAPL:Execute()

        if ShouldAoE() then
            AoEAPL:Execute()
        else
            DefaultAPL:Execute()
        end
    else
        OutOfCombatAPL:Execute()
    end
end)

-- Register the ElementalShaman module with Bastion
Bastion:Register(ElementalShamanModule)