local Tinkr, Bastion = ...
local AfflictionModule = Bastion.Module:New('AfflictionWarlock')
local Player = Bastion.UnitManager:Get('player')
local Target = Bastion.UnitManager:Get('target')
local SpellBook = Bastion.SpellBook:New()

-- Spells
local Agony = SpellBook:GetSpell(980)
local UnstableAffliction = SpellBook:GetSpell(316099)
local Corruption = SpellBook:GetSpell(172)
local MaleficRapture = SpellBook:GetSpell(324536)
local ShadowBolt = SpellBook:GetSpell(686)
local SeedOfCorruption = SpellBook:GetSpell(27243)
local DrainSoul = SpellBook:GetSpell(198590)
local Haunt = SpellBook:GetSpell(48181)
local VileTaint = SpellBook:GetSpell(278350)
local SoulRot = SpellBook:GetSpell(325640)
local SummonDarkglare = SpellBook:GetSpell(205180)
local PhantomSingularity = SpellBook:GetSpell(205179)
local SiphonLife = SpellBook:GetSpell(63106)

-- Create APLs
local DefaultAPL = Bastion.APL:New('default')
local AoEAPL = Bastion.APL:New('aoe')
local OpeningAPL = Bastion.APL:New('opening')
local CooldownAPL = Bastion.APL:New('cooldown')

-- Helper function to count active DoTs on target
local function CountDotsOnTarget(unit)
    local count = 0
    if unit:GetAuras():FindMy(Agony):IsUp() then count = count + 1 end
    if unit:GetAuras():FindMy(Corruption):IsUp() then count = count + 1 end
    if unit:GetAuras():FindMy(UnstableAffliction):IsUp() then count = count + 1 end
    if unit:GetAuras():FindMy(SiphonLife):IsUp() then count = count + 1 end
    if unit:GetAuras():FindMy(PhantomSingularity):IsUp() then count = count + 1 end
    if unit:GetAuras():FindMy(VileTaint):IsUp() then count = count + 1 end
    return count
end

-- Function to check if there are nearby enemies without DoTs
local function NearbyEnemiesNeedDots()
    local needDots = false
    Bastion.ObjectManager.enemies:each(function(unit)
        if unit:IsAlive() and Player:CanSee(unit) and Player:GetDistance(unit) <= 40 and unit:IsAffectingCombat() then
            if not unit:GetAuras():FindMy(Agony):IsUp() or not unit:GetAuras():FindMy(Corruption):IsUp() then
                needDots = true
                return true  -- Break the loop
            end
        end
    end)
    return needDots
end

-- Function to check if the target is valid for attacks
local function IsValidTarget(unit)
    return unit and unit:Exists() and unit:IsAlive() and unit:IsEnemy() and Player:CanSee(unit) and unit:IsAffectingCombat()
end

-- Default APL
DefaultAPL:AddSpell(
    Agony:CastableIf(function(self)
        return self:IsKnownAndUsable() and IsValidTarget(Target) and Target:GetAuras():FindMy(Agony):GetRemainingTime() < 5.4
    end):SetTarget(Target)
)

DefaultAPL:AddSpell(
    UnstableAffliction:CastableIf(function(self)
        return self:IsKnownAndUsable() and IsValidTarget(Target) and Target:GetAuras():FindMy(UnstableAffliction):GetRemainingTime() < 6.3
    end):SetTarget(Target)
)

DefaultAPL:AddSpell(
    Corruption:CastableIf(function(self)
        return self:IsKnownAndUsable() and IsValidTarget(Target) and Target:GetAuras():FindMy(Corruption):GetRemainingTime() < 4.2
    end):SetTarget(Target)
)

DefaultAPL:AddSpell(
    MaleficRapture:CastableIf(function(self)
        return self:IsKnownAndUsable() and IsValidTarget(Target) and 
               ((Player:GetPower() == Player:GetMaxPower() or Player:GetPower() >= 4) or NearbyEnemiesNeedDots())
    end):SetTarget(Target)
)

DefaultAPL:AddSpell(
    ShadowBolt:CastableIf(function(self)
        return self:IsKnownAndUsable() and IsValidTarget(Target) and Player:GetPower() < 4
    end):SetTarget(Target)
)

-- AoE APL
AoEAPL:AddSpell(
    SeedOfCorruption:CastableIf(function(self)
        return self:IsKnownAndUsable() and Bastion.ObjectManager.enemies:count() >= 3 and 
               IsValidTarget(Target) and not Target:GetAuras():FindMy(Corruption):IsUp()
    end):SetTarget(Target)
)

-- Opening APL
OpeningAPL:AddSpell(
    UnstableAffliction:CastableIf(function(self)
        return self:IsKnownAndUsable() and Target:Exists() and Target:IsEnemy() and not Player:IsAffectingCombat()
    end):SetTarget(Target)
)

OpeningAPL:AddSpell(
    Agony:CastableIf(function(self)
        return self:IsKnownAndUsable() and Target:Exists() and Target:IsEnemy() and not Player:IsAffectingCombat()
    end):SetTarget(Target)
)

OpeningAPL:AddSpell(
    Corruption:CastableIf(function(self)
        return self:IsKnownAndUsable() and Target:Exists() and Target:IsEnemy() and not Player:IsAffectingCombat()
    end):SetTarget(Target)
)

-- Cooldown APL
CooldownAPL:AddSpell(
    Haunt:CastableIf(function(self)
        return self:IsKnownAndUsable() and IsValidTarget(Target)
    end):SetTarget(Target)
)

CooldownAPL:AddSpell(
    VileTaint:CastableIf(function(self)
        return self:IsKnownAndUsable() and IsValidTarget(Target)
    end):SetTarget(Target):OnCast(function(self)
        if IsSpellPending() == 64 then
            local x, y, z = ObjectPosition(Target:GetOMToken())
            if x and y and z then
                self:Click(x, y, z)
            end
        end
    end)
)

CooldownAPL:AddSpell(
    SoulRot:CastableIf(function(self)
        return self:IsKnownAndUsable() and IsValidTarget(Target)
    end):SetTarget(Target)
)

CooldownAPL:AddSpell(
    SummonDarkglare:CastableIf(function(self)
        return self:IsKnownAndUsable() and IsValidTarget(Target) and CountDotsOnTarget(Target) >= 3
    end):SetTarget(Player):OnCast(function(self)
        if IsSpellPending() == 64 then
            local x, y, z = ObjectPosition(Player:GetOMToken())
            if x and y and z then
                self:Click(x, y, z)
            end
        end
    end)
)

AfflictionModule:Sync(function()
    if not Player:IsAffectingCombat() then
        if Target:Exists() and Target:IsEnemy() and Target:IsAlive() and not Target:IsAffectingCombat() then
            OpeningAPL:Execute()
        end
        return
    end

    if not IsValidTarget(Target) then
        return
    end

    CooldownAPL:Execute()

    if Bastion.ObjectManager.enemies:count() >= 3 then
        AoEAPL:Execute()
    end

    DefaultAPL:Execute()
end)

Bastion:Register(AfflictionModule)