local Tinkr, Bastion = ...

local RestoShamanModule = Bastion.Module:New('RestoShaman')
local Player = Bastion.UnitManager:Get('player')
local Target = Bastion.UnitManager:Get('target')
local SpellBook = Bastion.SpellBook:New()

-- Spells
local WaterShield = SpellBook:GetSpell(52127)
local EarthShield = SpellBook:GetSpell(974)
local Riptide = SpellBook:GetSpell(61295)
local HealingSurge = SpellBook:GetSpell(8004)
local HealingWave = SpellBook:GetSpell(77472)
local ChainHeal = SpellBook:GetSpell(1064)
local HealingRain = SpellBook:GetSpell(73920)
local Downpour = SpellBook:GetSpell(207778)
local HealingStreamTotem = SpellBook:GetSpell(5394)
local CloudburstTotem = SpellBook:GetSpell(157153)
local SpiritLinkTotem = SpellBook:GetSpell(98008)
local AncestralGuidance = SpellBook:GetSpell(108281)
local Ascendance = SpellBook:GetSpell(114052)
local Wellspring = SpellBook:GetSpell(197995)
local UnleashLife = SpellBook:GetSpell(73685)
local PrimordialWave = SpellBook:GetSpell(375982)
local LavaBurst = SpellBook:GetSpell(51505)
local FlameShock = SpellBook:GetSpell(188389)
local WindShear = SpellBook:GetSpell(57994)
local LightningBolt = SpellBook:GetSpell(403)

-- Buffs
local TidalWaves = SpellBook:GetSpell(53390)
local WaterShieldBuff = SpellBook:GetSpell(52127)
local EarthShieldBuff = SpellBook:GetSpell(383648)
local UnleashLifeBuff = SpellBook:GetSpell(73685)

-- Custom Units
local Lowest = Bastion.UnitManager:CreateCustomUnit('lowest', function()
    local lowest = nil
    local lowestHP = 100

    Bastion.UnitManager:EnumFriends(function(unit)
        if unit:IsDead() or Player:GetDistance(unit) > 40 or not Player:CanSee(unit) then
            return false
        end

        local hp = unit:GetHP()
        if hp < lowestHP then
            lowest = unit
            lowestHP = hp
        end
    end)

    return lowest or Player
end)

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 CooldownAPL = Bastion.APL:New('cooldown')
local DefensiveAPL = Bastion.APL:New('defensive')
local DpsAPL = Bastion.APL:New('dps')
local InterruptAPL = Bastion.APL:New('interrupt')

-- Helper Functions
local function GetPlayerManaPercent()
    return (UnitPower("player", Enum.PowerType.Mana) / UnitPowerMax("player", Enum.PowerType.Mana)) * 100
end

local function NeedsUrgentHealing()
    return Lowest:GetHP() < 70 or Player:GetPartyHPAround(30, 80) >= 3
end

-- Modify the EarthShield spell in the DefaultAPL
DefaultAPL:AddSpell(
    EarthShield:CastableIf(function(self)
        local playerShield = Player:GetAuras():FindMy(EarthShieldBuff)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling() and
            (not playerShield:IsUp() or playerShield:GetRemainingTime() < 30)
    end):SetTarget(Player)
)

-- Default APL
DefaultAPL:AddSpell(
    Riptide:CastableIf(function(self)
        return Lowest:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and Lowest:GetHP() < 95 and not Lowest:GetAuras():FindMy(TidalWaves):IsUp()
    end):SetTarget(Lowest)
)

DefaultAPL:AddSpell(
    UnleashLife:CastableIf(function(self)
        return Lowest:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and Lowest:GetHP() < 80
    end):SetTarget(Lowest)
)

DefaultAPL:AddSpell(
    HealingSurge:CastableIf(function(self)
        return Lowest:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and Lowest:GetHP() < 70 and Player:GetAuras():FindMy(TidalWaves):IsUp()
    end):SetTarget(Lowest)
)

DefaultAPL:AddSpell(
    HealingWave:CastableIf(function(self)
        return Lowest:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and Lowest:GetHP() < 85
    end):SetTarget(Lowest)
)

DefaultAPL:AddSpell(
    ChainHeal:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and Player:GetPartyHPAround(40, 85) >= 3
    end):SetTarget(Lowest)
)

-- Cooldown APL
CooldownAPL:AddSpell(
    CloudburstTotem:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and Player:GetPartyHPAround(40, 80) >= 3
    end):SetTarget(Player)
)

CooldownAPL:AddSpell(
    HealingStreamTotem:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and not CloudburstTotem:IsKnownAndUsable() and Player:GetPartyHPAround(40, 90) >= 2
    end):SetTarget(Player)
)

CooldownAPL:AddSpell(
    SpiritLinkTotem:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and Player:GetPartyHPAround(30, 70) >= 3
    end):SetTarget(Player)
)

CooldownAPL:AddSpell(
    AncestralGuidance:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and Player:GetPartyHPAround(40, 75) >= 3
    end):SetTarget(Player)
)

CooldownAPL:AddSpell(
    Ascendance:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and Player:GetPartyHPAround(40, 60) >= 3
    end):SetTarget(Player)
)

-- DPS APL
DpsAPL:AddSpell(
    FlameShock:CastableIf(function(self)
        return BestTarget:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and not BestTarget:GetAuras():FindMy(FlameShock):IsUp()
    end):SetTarget(BestTarget)
)

DpsAPL:AddSpell(
    LavaBurst:CastableIf(function(self)
        return BestTarget:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and BestTarget:GetAuras():FindMy(FlameShock):IsUp()
    end):SetTarget(BestTarget)
)

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

-- Interrupt APL
InterruptAPL:AddSpell(
    WindShear:CastableIf(function(self)
        return BestTarget:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and BestTarget:IsCasting() and BestTarget:IsInterruptible()
            and BestTarget:GetChannelOrCastPercentComplete() >= 50
    end):SetTarget(BestTarget)
)

-- Modify the out of combat Earth Shield logic in the Module Sync function
RestoShamanModule:Sync(function()
    if Player:IsMounted() then
        return
    end

    -- Out of combat healing and buffs
    if not Player:IsAffectingCombat() then
        if not Player:GetAuras():FindMy(WaterShieldBuff):IsUp() then
            WaterShield:Cast(Player)
        elseif not Player:GetAuras():FindMy(EarthShieldBuff):IsUp() then
            EarthShield:Cast(Player)
        elseif Lowest:GetHP() < 90 then
            DefaultAPL:Execute()
        end
        return
    end

    -- In combat
    InterruptAPL:Execute()

    if NeedsUrgentHealing() then
        CooldownAPL:Execute()
        DefaultAPL:Execute()
    else
        if GetPlayerManaPercent() > 80 and BestTarget:Exists() then
            DpsAPL:Execute()
        end
        
        if Riptide:GetCharges() >= 2 then
            Bastion.UnitManager:EnumFriends(function(unit)
                if unit:GetHP() < 95 and not unit:GetAuras():FindMy(TidalWaves):IsUp() then
                    Riptide:Cast(unit)
                    return true  -- Stop enumerating after casting
                end
            end)
        end
        
        if HealingWave:IsKnownAndUsable() and Lowest:GetHP() < 90 then
            HealingWave:Cast(Lowest)
        end
    end

    -- AoE Healing
    if Player:GetPartyHPAround(30, 85) >= 4 then
        if HealingRain:IsKnownAndUsable() then
            local loc = Bastion.UnitManager:FindFriendsCentroid(10, 40)
            if loc then
                HealingRain:Cast(Player):OnCast(function(self)
                    self:Click(loc)
                end)
            end
        elseif Downpour:IsKnownAndUsable() then
            local loc = Bastion.UnitManager:FindFriendsCentroid(10, 40)
            if loc then
                Downpour:Cast(Player):OnCast(function(self)
                    self:Click(loc)
                end)
            end
        elseif Wellspring:IsKnownAndUsable() then
            Wellspring:Cast(Player)
        end
    end
end)

Bastion:Register(RestoShamanModule)