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)
local LightningShield = SpellBook:GetSpell(192106)
local EarthShield = SpellBook:GetSpell(974)

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

local InterruptTarget = Bastion.UnitManager:CreateCustomUnit('interrupttarget', function()
    local target = nil
    Bastion.UnitManager:EnumEnemies(function(unit)
        if unit:IsAffectingCombat() and unit:GetDistance(Player) <= 30 and Player:CanSee(unit)
           and unit:IsCasting() and unit:IsInterruptible() then
            local castPercentage = unit:GetChannelOrCastPercentComplete()
            if castPercentage >= 30 and castPercentage <= 80 then
                target = unit
                return true
            end
        end
    end)
    return target or Bastion.UnitManager:Get('none')
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')
local InterruptAPL = Bastion.APL:New('interrupt')

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

local function NeedSkyfury()
    -- Check if any party member (including the player) needs Skyfury
    local needsBuff = false
    Bastion.UnitManager:EnumFriends(function(unit)
        if unit:IsAlive() and not unit:GetAuras():FindAny(Skyfury):IsUp() then
            needsBuff = true
            return true  -- This will break the enumeration loop
        end
    end)
    return needsBuff
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)
)

-- OutOfCombat APL
OutOfCombatAPL:AddSpell(
    Skyfury:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling() and NeedSkyfury()
    end):SetTarget(Player)
)

OutOfCombatAPL:AddSpell(
    EarthShield:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling() and not Player:GetAuras():FindAny(EarthShield_Buff):IsUp()
    end):SetTarget(Player)
)

OutOfCombatAPL:AddSpell(
    LightningShield:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling() and not Player:GetAuras():FindAny(LightningShield_Buff):IsUp()
    end):SetTarget(Player)
)

-- 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, 39)
        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
InterruptAPL:AddSpell(
    WindShear:CastableIf(function(self)
        local target = InterruptTarget
        local randomInterruptPercent = math.random(50, 90)
        return target:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and target:IsCasting() and target:IsInterruptible()
            and target:GetChannelOrCastPercentComplete() >= randomInterruptPercent
    end):SetTarget(InterruptTarget)
)

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

    if Player:IsAffectingCombat() then
        DefensiveAPL:Execute()
        InterruptAPL:Execute()
        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)