local Tinkr, Bastion = ...

local AugmentationModule = Bastion.Module:New('AugmentationEvoker')
local Player = Bastion.UnitManager:Get('player')
local Target = Bastion.UnitManager:Get('target')
local Focus = Bastion.UnitManager:Get('focus')
local SpellBook = Bastion.SpellBook:New()

-- Ensure EventManager is properly initialized
local EventManager = Bastion.EventManager:New()

-- Constants for Ebon Might logic
local EBON_MIGHT_MIN_ENEMIES = 1
local EBON_MIGHT_MIN_TTD = 12 -- seconds
local EBON_MIGHT_MIN_ESSENCE = 50 -- percent of max essence

-- Spells
local Prescience = SpellBook:GetSpell(409311)
local EbonMight = SpellBook:GetSpell(395152)
local FireBreath = SpellBook:GetSpell(382266)
local Upheaval = SpellBook:GetSpell(396286)
local Eruption = SpellBook:GetSpell(395160)
local LivingFlame = SpellBook:GetSpell(361469)
local TipTheScales = SpellBook:GetSpell(370553)
local Hover = SpellBook:GetSpell(358267)
local AzureStrike = SpellBook:GetSpell(362969)
local BlisteringScales = SpellBook:GetSpell(360827)
local TimeSkip = SpellBook:GetSpell(404977)

-- Buffs and Debuffs
local PrescienceBuff = SpellBook:GetSpell(410089)
local EbonMightBuff = SpellBook:GetSpell(395296)
local EssenceBurstBuff = SpellBook:GetSpell(369299)

-- Talents
local LeapingFlames = SpellBook:GetSpell(370901)

-- Create APLs
local CooldownAPL = Bastion.APL:New('cooldown')
local BuffAPL = Bastion.APL:New('buff')
local DamageAPL = Bastion.APL:New('damage')
local MovementAPL = Bastion.APL:New('movement')

-- Empowered spell handling
local empowering = {}

function Player:GetEmpoweredStage()
    local stage = 0
    local _, _, _, startTime, _, _, _, spellID, _, numStages = UnitChannelInfo(self:GetOMToken())

    if numStages and numStages > 0 then
        startTime = startTime / 1000
        local currentTime = GetTime()
        local stageDuration = 0
        for i = 1, numStages do
            stageDuration = stageDuration + GetUnitEmpowerStageDuration((self:GetOMToken()), i - 1) / 1000
            if startTime + stageDuration > currentTime then
                break
            end
            stage = i
        end
    end
    return stage
end

EventManager:RegisterWoWEvent("UNIT_SPELLCAST_EMPOWER_START", function(...)
    local unit, _, id = ...
    if not unit then return end
    local guid = UnitGUID(unit)
    if not guid then return end
    empowering[guid] = -1
end)

EventManager:RegisterWoWEvent("UNIT_SPELLCAST_EMPOWER_STOP", function(...)
    local unit, _, id = ...
    if not unit then return end
    local guid = UnitGUID(unit)
    if not guid then return end
    empowering[guid] = nil
end)

function Player:GetUnitEmpowerStage()
    local name, _, _, startTime, endTime, _, _, _, _, numStages = UnitChannelInfo(self:GetOMToken())

    if name and empowering[self:GetGUID()] == -1 then
        empowering[self:GetGUID()] = numStages
    end

    if not name then
        return empowering[self:GetGUID()] or 0
    end

    local getStageDuration = function(stage)
        if stage == numStages then
            return GetUnitEmpowerHoldAtMaxTime(self:GetOMToken())
        else
            return GetUnitEmpowerStageDuration(self:GetOMToken(), stage - 1)
        end
    end

    local time = GetTime() - (startTime / 1000)

    local higheststage = 0
    local sumdur = 0
    for i = 1, numStages - 1 do
        local duration = getStageDuration(i) / 1000
        sumdur = sumdur + duration

        if time > sumdur then
            higheststage = i
        end
    end

    return higheststage
end

-- Helper Functions
local function GetPrescienceTargets()
    local targets = Bastion.List:New()
    if Focus:Exists() and not Focus:IsTank() and Player:GetDistance(Focus) <= 40 and Player:CanSee(Focus) then
        targets:push(Focus)
    end
    Bastion.UnitManager:EnumFriends(function(unit)
        if targets:count() >= 2 then return true end
        if unit ~= Focus and unit:IsPlayer() and not unit:IsTank() and Player:GetDistance(unit) <= 40 and Player:CanSee(unit) then
            local aura = unit:GetAuras():FindMy(PrescienceBuff)
            if not aura:IsUp() or aura:GetRemainingTime() < 10 then
                targets:push(unit)
            end
        end
    end)
    return targets
end

local function IsEbonMightActive()
    return Player:GetAuras():FindMy(EbonMightBuff):IsUp()
end

local function IsEssenceCapped()
    return Player:GetPower() >= Player:GetMaxPower() or Player:GetAuras():FindMy(EssenceBurstBuff):GetCount() == 2
end

local function ShouldCastEbonMight()
    local viableEnemyCount = 0
    local playerEssencePercent = (Player:GetPower() / Player:GetMaxPower()) * 100

    Bastion.UnitManager:EnumEnemies(function(unit)
        if unit:IsAffectingCombat() and unit:TimeToDie() >= EBON_MIGHT_MIN_TTD then
            viableEnemyCount = viableEnemyCount + 1
        end
    end)

    local isBossFight = Target:Exists() and Target:IsBoss()
    local hasEnoughEssence = playerEssencePercent >= EBON_MIGHT_MIN_ESSENCE

    if isBossFight then
        return hasEnoughEssence and Target:TimeToDie() >= EBON_MIGHT_MIN_TTD
    else
        return viableEnemyCount >= EBON_MIGHT_MIN_ENEMIES and hasEnoughEssence
    end
end

-- Custom Units
local PrescienceTarget = Bastion.UnitManager:CreateCustomUnit('presciencetarget', function()
    local targets = GetPrescienceTargets()
    return targets:count() > 0 and targets:peek() or Player
end)

local EbonMightTarget = Bastion.UnitManager:CreateCustomUnit('ebonmighttarget', function()
    local target = nil
    Bastion.UnitManager:EnumFriends(function(unit)
        if unit:IsPlayer() and Player:GetDistance(unit) <= 40 and Player:CanSee(unit) then
            target = unit
            return true
        end
    end)
    return target or Player
end)

local DamageTarget = Bastion.UnitManager:CreateCustomUnit('damagetarget', function()
    local target = nil
    local highestHealth = 0
    Bastion.UnitManager:EnumEnemies(function(unit)
        if unit:IsAffectingCombat() and Player:IsWithinCombatDistance(unit, 40) and Player:CanSee(unit) then
            local health = unit:GetHealth()
            if health > highestHealth then
                target = unit
                highestHealth = health
            end
        end
    end)
    return target or Target
end)

local TankTarget = Bastion.UnitManager:CreateCustomUnit('tanktarget', function()
    local tank = nil
    Bastion.UnitManager:EnumFriends(function(unit)
        if unit:IsTank() and Player:GetDistance(unit) <= 40 and Player:CanSee(unit) then
            tank = unit
            return true
        end
    end)
    return tank or Player
end)

-- Cooldown APL
CooldownAPL:AddSpell(
    EbonMight:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling() and ShouldCastEbonMight()
    end):SetTarget(EbonMightTarget)
)

CooldownAPL:AddSpell(
    TipTheScales:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling() and
               IsEbonMightActive() and (FireBreath:IsKnownAndUsable() or Upheaval:IsKnownAndUsable())
    end):SetTarget(Player)
)

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

-- Buff APL
BuffAPL:AddSpell(
    Prescience:CastableIf(function(self)
        return GetPrescienceTargets():count() > 0 and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling()
    end):SetTarget(PrescienceTarget):OnCast(function(self)
        local targets = GetPrescienceTargets()
        if targets:count() > 1 then
            self:Cast(targets[2] or Player)
        end
    end)
)

BuffAPL:AddSpell(
    BlisteringScales:CastableIf(function(self)
        return TankTarget:Exists() and self:IsKnownAndUsable() and not Player:IsCastingOrChanneling() and
               not TankTarget:GetAuras():FindMy(BlisteringScales):IsUp()
    end):SetTarget(TankTarget)
)

-- Damage APL
DamageAPL:AddSpell(
    FireBreath:CastableIf(function(self)
        return DamageTarget:IsAffectingCombat() and IsEbonMightActive()
    end):SetTarget(DamageTarget):OnCast(function(self)
        local empowerStage = Player:GetUnitEmpowerStage()
        if empowerStage >= 1 then
            C_Timer.After(0.1, function()
                self:ForceCast(DamageTarget)
            end)
        end
    end)
)

DamageAPL:AddSpell(
    Upheaval:CastableIf(function(self)
        return DamageTarget:IsAffectingCombat() and IsEbonMightActive()
    end):SetTarget(DamageTarget):OnCast(function(self)
        local empowerStage = Player:GetUnitEmpowerStage()
        if empowerStage >= 1 then
            self:ForceCast(DamageTarget)
        end
    end)
)

DamageAPL:AddSpell(
    Eruption:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling() and
               (IsEbonMightActive() or IsEssenceCapped())
    end):SetTarget(DamageTarget)
)

DamageAPL:AddSpell(
    LivingFlame:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:IsCastingOrChanneling() and
               (LeapingFlames:IsKnown() or not FireBreath:IsKnownAndUsable())
    end):SetTarget(DamageTarget)
)

-- Movement APL
MovementAPL:AddSpell(
    Hover:CastableIf(function(self)
        return self:IsKnownAndUsable() and not Player:GetAuras():FindMy(Hover):IsUp()
    end):SetTarget(Player)
)

MovementAPL:AddSpell(
    AzureStrike:CastableIf(function(self)
        return self:IsKnownAndUsable() and Player:IsMoving() and not Player:GetAuras():FindMy(Hover):IsUp()
    end):SetTarget(DamageTarget)
)

-- Module Sync
AugmentationModule:Sync(function()
    -- Check if the player is mounted
    if Player:IsMounted() then
        return  -- Exit the function if mounted, effectively disabling all APLs
    end

    if not Player:IsAffectingCombat() then
        BuffAPL:Execute()
        return
    end

    CooldownAPL:Execute()
    BuffAPL:Execute()

    if DamageAPL:Execute() then
        return
    end

    MovementAPL:Execute()
end)

Bastion:Register(AugmentationModule)