local Tinkr, Bastion = ...

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

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

-- Spells
local WaterShield = SpellBook:GetSpell(52127)
local EarthShield = SpellBook:GetSpell(974)
local HealingRain = SpellBook:GetSpell(73920)
local Riptide = SpellBook:GetSpell(61295)
local CloudburstTotem = SpellBook:GetSpell(157153)
local HealingStreamTotem = SpellBook:GetSpell(5394)
local PrimordialWave = SpellBook:GetSpell(375982)
local UnleashLife = SpellBook:GetSpell(73685)
local Downpour = SpellBook:GetSpell(207778)
local ChainHeal = SpellBook:GetSpell(1064)
local HealingWave = SpellBook:GetSpell(77472)
local HealingSurge = SpellBook:GetSpell(8004)
local Wellspring = SpellBook:GetSpell(197995)
local SpiritLinkTotem = SpellBook:GetSpell(98008)
local HealingTideTotem = SpellBook:GetSpell(108280)
local AncestralGuidance = SpellBook:GetSpell(108281)
local ManaTideTotem = SpellBook:GetSpell(16191)
local Ascendance = SpellBook:GetSpell(114052)
local NaturesSwiftness = SpellBook:GetSpell(378081)
local EarthenWallTotem = SpellBook:GetSpell(198838)
local PurifySpirit = SpellBook:GetSpell(77130)
local LightningBolt = SpellBook:GetSpell(403)
local ChainLightning = SpellBook:GetSpell(421)
local FlameShock = SpellBook:GetSpell(188389)
local LavaBurst = SpellBook:GetSpell(51505)

-- Buffs
local TidalWaves = SpellBook:GetSpell(53390)
local HighTide = SpellBook:GetSpell(288675)

-- Totem timers
local cloudburstTotemEnd = 0
local healingStreamTotemEnd = 0
local earthenWallTotemEnd = 0
local spiritLinkTotemEnd = 0
local healingTideTotemEnd = 0

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

    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)

-- APLs
local PrePullAPL = Bastion.APL:New('prepull')
local DefaultAPL = Bastion.APL:New('default')
local CooldownAPL = Bastion.APL:New('cooldown')
local DamageAPL = Bastion.APL:New('damage')

-- Debug function
local function Debug(message)
    -- print("[RestoShaman Debug]: " .. message)
end

-- Helper Functions
local function IsTotemActive(totemEnd)
    return GetTime() < totemEnd
end

local function UpdateTotemTimer(totemEnd, duration)
    return GetTime() + duration
end

local function ShouldUseCloudburstTotem()
    return not IsTotemActive(cloudburstTotemEnd) and (Player:GetEnemies(10) >= 3 or Player:GetPartyHPAround(30, 80) >= 3)
end

local function ShouldUseHealingStreamTotem()
    return not IsTotemActive(healingStreamTotemEnd) and not CloudburstTotem:IsKnownAndUsable()
end

local function ShouldUseEarthenWallTotem()
    return not IsTotemActive(earthenWallTotemEnd) and Player:GetPartyHPAround(40, 80) >= 3
end

local function ShouldUseSpiritLinkTotem()
    return not IsTotemActive(spiritLinkTotemEnd) and Player:GetPartyHPAround(40, 60) >= 3
end

local function ShouldUseHealingTideTotem()
    return not IsTotemActive(healingTideTotemEnd) and Player:GetPartyHPAround(40, 70) >= 3
end

local function GetRiptideCount()
    local count = 0
    Bastion.UnitManager:EnumFriends(function(unit)
        if unit:GetAuras():FindMy(Riptide):IsUp() then
            count = count + 1
        end
    end)
    return count
end

local function ShouldUseAscendance()
    return Player:GetPartyHPAround(40, 70) >= 3 and CloudburstTotem:GetTimeSinceLastCast() < 3
end

local function GetTanks()
    local tanks = {}
    Bastion.UnitManager:EnumFriends(function(unit)
        if unit:IsTank() and not unit:IsDead() and Player:CanSee(unit) and Player:GetDistance(unit) <= 40 then
            table.insert(tanks, unit)
        end
    end)
    Debug("Found " .. #tanks .. " tanks")
    return tanks
end

local function ApplyEarthShield()
    local tanks = GetTanks()
    for _, tank in ipairs(tanks) do
        if not tank:GetAuras():FindMy(EarthShield):IsUp() then
            Debug("Applying Earth Shield to " .. tank:GetName())
            return EarthShield:Cast(tank)
        end
    end
    Debug("No tanks need Earth Shield")
end

local function NeedsHealing(unit, threshold)
    return unit:GetHP() < threshold
end

-- Pre-Pull APL
PrePullAPL:AddSpell(
    WaterShield:CastableIf(function(self)
        return Player:Exists() and self:IsKnownAndUsable() and not Player:GetAuras():FindMy(WaterShield):IsUp()
    end):SetTarget(Player)
)

-- Default APL
DefaultAPL:AddSpell(
    CloudburstTotem:CastableIf(function(self)
        return Player:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling() and ShouldUseCloudburstTotem()
    end):SetTarget(Player):OnCast(function()
        cloudburstTotemEnd = UpdateTotemTimer(cloudburstTotemEnd, 15)
    end)
)

DefaultAPL:AddSpell(
    HealingStreamTotem:CastableIf(function(self)
        return Player:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling() and ShouldUseHealingStreamTotem()
    end):SetTarget(Player):OnCast(function()
        healingStreamTotemEnd = UpdateTotemTimer(healingStreamTotemEnd, 15)
    end)
)

DefaultAPL:AddSpell(
    Riptide:CastableIf(function(self)
        return Lowest:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and NeedsHealing(Lowest, 90)
    end):SetTarget(Lowest)
)

DefaultAPL:AddSpell(
    PrimordialWave:CastableIf(function(self)
        return Lowest:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and NeedsHealing(Lowest, 85)
    end):SetTarget(Lowest)
)

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

DefaultAPL:AddSpell(
    HealingRain:CastableIf(function(self)
        return Player:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and Player:GetPartyHPAround(30, 85) >= 6 and not Player:IsMoving()
    end):SetTarget(Player)
)

DefaultAPL:AddSpell(
    Downpour:CastableIf(function(self)
        return Player:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and Player:GetPartyHPAround(30, 80) >= 3
    end):SetTarget(Player)
)

DefaultAPL:AddSpell(
    ChainHeal:CastableIf(function(self)
        return Lowest:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and (Player:GetAuras():FindMy(HighTide):IsUp() or Player:GetAuras():FindMy(NaturesSwiftness):IsUp())
            and not Player:IsMoving() and NeedsHealing(Lowest, 75)
    end):SetTarget(Lowest)
)

DefaultAPL:AddSpell(
    HealingWave:CastableIf(function(self)
        return Lowest:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and Player:GetAuras():FindMy(TidalWaves):IsUp() and not Player:IsMoving()
            and NeedsHealing(Lowest, 85)
    end):SetTarget(Lowest)
)

DefaultAPL:AddSpell(
    HealingSurge:CastableIf(function(self)
        return Lowest:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and NeedsHealing(Lowest, 70)
    end):SetTarget(Lowest)
)

DefaultAPL:AddSpell(
    PurifySpirit:CastableIf(function(self)
        return Lowest:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and Lowest:GetAuras():HasAnyDispelableAura(PurifySpirit)
    end):SetTarget(Lowest)
)

-- Cooldown APL
CooldownAPL:AddSpell(
    SpiritLinkTotem:CastableIf(function(self)
        return Player:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling() and ShouldUseSpiritLinkTotem()
    end):SetTarget(Player):OnCast(function()
        spiritLinkTotemEnd = UpdateTotemTimer(spiritLinkTotemEnd, 6)
    end)
)

CooldownAPL:AddSpell(
    HealingTideTotem:CastableIf(function(self)
        return Player:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling() and ShouldUseHealingTideTotem()
    end):SetTarget(Player):OnCast(function()
        healingTideTotemEnd = UpdateTotemTimer(healingTideTotemEnd, 10)
    end)
)

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

CooldownAPL:AddSpell(
    ManaTideTotem:CastableIf(function(self)
        return Player:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and Player:GetPP() <= 80
    end):SetTarget(Player)
)

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

CooldownAPL:AddSpell(
    NaturesSwiftness:CastableIf(function(self)
        return Lowest:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and NeedsHealing(Lowest, 50)
    end):SetTarget(Lowest)
)

CooldownAPL:AddSpell(
    EarthenWallTotem:CastableIf(function(self)
        return Player:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling() and ShouldUseEarthenWallTotem()
    end):SetTarget(Player):OnCast(function()
        earthenWallTotemEnd = UpdateTotemTimer(earthenWallTotemEnd, 15)
    end)
)

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

-- Damage APL
DamageAPL:AddSpell(
    FlameShock:CastableIf(function(self)
        return Target:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and not Target:GetAuras():FindMy(FlameShock):IsUp()
    end):SetTarget(Target)
)

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

DamageAPL:AddSpell(
    ChainLightning:CastableIf(function(self)
        return Target:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
            and Player:GetEnemies(10) > 2
    end):SetTarget(Target)
)

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

-- Module Sync
RestoShamanModule:Sync(function()
    if not Player:IsAffectingCombat() then
        -- Reset totem timers when leaving combat
        cloudburstTotemEnd = 0
        healingStreamTotemEnd = 0
        earthenWallTotemEnd = 0
        spiritLinkTotemEnd = 0
        healingTideTotemEnd = 0

        if not Player:GetAuras():FindMy(WaterShield):IsUp() then
            WaterShield:Cast(Player)
        end
        ApplyEarthShield()
    elseif Player:IsAffectingCombat() then
        ApplyEarthShield()
        CooldownAPL:Execute()
        DefaultAPL:Execute()
        
        -- Use damage abilities when healing isn't needed
        if Player:GetPP() > 80 and Player:GetPartyHPAround(40, 95) == 0 then
            DamageAPL:Execute()
        end
    end
end)

Bastion:Register(RestoShamanModule)