ally system

--!optimize 2
--!native

local RepStorage = game:FindService("ReplicatedStorage")
local ServerStorage = game:FindService("ServerStorage")
local RunService = game:GetService("RunService")

local remotes = RepStorage.Remotes

local clientModules = RepStorage.Modules
local allyData = require(clientModules.AllyData)
local pathModule = require(clientModules.PathModule)

local serverModules = ServerStorage.Modules
local enemyHandler = require(serverModules.EnemyHandler)
local targetingFuncs = require(serverModules.TargetingFunctions)

local DamageEnemy = enemyHandler.DamageEnemy
local RemoveEnemy = enemyHandler.RemoveEnemy
local FindNodeAfterPoint = pathModule.FindNodeAfterPoint
local FindNodeBeforePoint = pathModule.FindNodeBeforePoint
local paths = pathModule.paths
local enemies = enemyHandler.enemies

local IDcount = 0
local allies = {}

--

local IDX_POS = 1
local IDX_DIST = 2

--

local function RemoveAlly(ally)
    ally.dead = true

    local ID = ally.ID

    allies[ID] = nil
    remotes.RemoveAlly:FireAllClients(ID)

    local DeathFunc = ally.funcs.OnDeath
    if DeathFunc then
        task.spawn(DeathFunc, ally)
    end
end

local function DamageAlly(ally, damageAmount: number, bypassShield: boolean?) --returns true damage amount
    local health = ally.health
    local shield = ally.shield

    if shield and shield ~= 0 and not bypassShield then
        local allyID = ally.ID
        local buf = buffer.create(8)
        buffer.writeu32(buf, 0, allyID)

        if damageAmount >= shield then
            ally.shield = 0
            local ShieldBreakFunc = ally.funcs.OnShieldBreak
            if ShieldBreakFunc then task.spawn(buf) end
            remotes.UpdateAllyShield:FireAllClients(buf)

            return shield
        end

        local newShield = shield - damageAmount
        ally.shield = newShield

        buffer.writef32(buf, 4, newShield)
        remotes.UpdateAllyShield:FireAllClients(buf)
    else
        if damageAmount >= health then
            RemoveAlly(ally)
            return health
        end

        local newHealth = health - damageAmount
        ally.health = newHealth

        local allyID = ally.ID
        local buf = buffer.create(8)
        buffer.writeu32(buf, 0, allyID)
        buffer.writef32(buf, 4, newHealth)
        remotes.UpdateAllyHealth:FireAllClients(buf)
    end

    return damageAmount
end

local function UnitRam(ally, enemy)
    local allyShield = ally.shield
    local allyHealth = ally.health
    DamageEnemy(enemy, if allyShield then allyHealth + allyShield else allyHealth, ally.player_name, nil, true)

    local enemyShield = enemy.shield
    local enemyHealth = enemy.health
    DamageAlly(ally, if enemyShield then enemyHealth + enemyShield else enemyHealth)
    return ally.dead
end

local function AddAlly(spawnData)
    IDcount += 1

    local allyType = spawnData.kind
    local playerName = spawnData.player
    local distance = spawnData.distance
    local pathIndex = spawnData.path or math.random(#paths)

    local allyInfo = allyData[allyType]
    local path = paths[pathIndex]
    local pathLength = #path
    local funcs = allyInfo.Functions and require(allyInfo.Functions)

    local offset = Vector3.new(
        if allyInfo.NoPathOffset then 0 else Random.new(IDcount):NextNumber(-0.5, 0.5),
        allyInfo.Model:GetExtentsSize().Y / 2,
        0
    )

    --

    remotes.AddAlly:FireAllClients(IDcount, workspace:GetServerTimeNow(), allyType, distance, pathIndex)

    local nodeIndex, trackPos, trackCFrame
    if distance then
        nodeIndex = FindNodeBeforePoint(path, pathLength, distance)

        local nodeA = path[nodeIndex]
        local nodeB = path[nodeIndex+1]

        local dir = (nodeB[IDX_POS] - nodeA[IDX_POS]).Unit
        trackPos = nodeA[IDX_POS] + dir * (distance - nodeA[IDX_DIST])
        trackCFrame = CFrame.lookAlong(trackPos, -dir)
    else
        nodeIndex = pathLength
        trackPos = path[pathLength][IDX_POS]
        trackCFrame = CFrame.lookAt(trackPos, path[pathLength-1][IDX_POS])
    end

    if funcs and funcs.Main then
        task.defer(funcs.Main, IDcount)
    end

    local enemy = {
        ID = IDcount,
        ally_type = allyType,
        path = pathIndex,
        player_name = playerName,

        track_pos = trackPos,
        true_pos = trackCFrame * offset,
        offset_normal = offset,
        distance = distance or path[pathLength][IDX_DIST],
        dead = false,

        display_speed = nil,
        speed = allyInfo.Speed,
        health = allyInfo.MaxHealth,
        max_health = allyInfo.MaxHealth,
        shield = allyInfo.Shield,
        ghost = allyInfo.Ghost,
        float = allyInfo.Float,

        node_index = nodeIndex,
        funcs = funcs or {}
    }

    allies[IDcount] = enemy
    return enemy
end

local function StepAlly(ally, amount: number)
    local backwards = amount < 0

    local path = paths[ally.path]
    local newDist = math.min(ally.distance - amount, path[#path][IDX_DIST])

    if newDist > 0 then
        local nodeIndex = (if backwards then FindNodeAfterPoint else FindNodeBeforePoint)(path, ally.node_index, newDist)

        local nodeA = path[nodeIndex]
        local nodeB = path[nodeIndex+1]

        local dir = (nodeB[IDX_POS] - nodeA[IDX_POS]).Unit
        local allyTrackPos = (nodeA[IDX_POS] + dir * (newDist - nodeA[IDX_DIST]))

        --

        if not ally.ghost then
            local allyFloat = ally.float
            local allyPos = if allyFloat then allyTrackPos + (Vector3.yAxis * allyFloat) else allyTrackPos
            local rangeSquared = allyData[ally.ally_type].Hit_Range^2

            for _, enemy in enemies do
                local dif = (enemy.track_pos - allyPos)
                local enemyFloat = enemy.float

                if (dif.X^2 + (if enemyFloat then dif.Y + enemyFloat else dif.Y)^2 + dif.Z^2) <= rangeSquared and UnitRam(ally, enemy) then
                    return
                end
            end
        end

        ally.distance = newDist
        ally.node_index = nodeIndex
        ally.track_pos = allyTrackPos
        ally.true_pos = CFrame.lookAlong(allyTrackPos, dir) * ally.offset_normal
    else
        RemoveAlly(ally)
    end
end

--

RunService.PreSimulation:Connect(function(dt)
    for _, ally in allies do
        StepAlly(ally, (ally.display_speed or ally.speed) * dt)
    end
end)

--

local function GetAllyCount()
    local count = 0

    for _ in allies do
        count += 1
    end
    return count
end

local function GetRandomAlly(flyingDet: boolean?)
    local indicies = {}
    local count = 0

    for ID, ally in allies do
        if not flyingDet and ally.float then continue end

        count += 1
        indicies[count] = ID
    end
    return #indicies ~= 0 and indicies[math.random(#indicies)]
end

local function GetAlliesInRange(pos: Vector3, range: number, flyingDet: boolean?)
    local rangeSquared = range^2

    local validIndicies = {}
    local count = 0

    for ID, ally in allies do
        if not flyingDet and ally.float then continue end

        local dif = (allies.track_pos - pos)
        local float = ally.Float

        if (dif.X^2 + (if float then dif.Y + float else dif.Y)^2 + dif.Z^2) <= rangeSquared then
            count += 1
            validIndicies[count] = ID
        end
    end
    return #validIndicies ~= 0 and validIndicies
end

local function GetAlliesInRange2D(pos: Vector3, range: number, flyingDet: boolean?)
    local rangeSquared = range^2

    local validIndicies = {}
    local count = 0

    for ID, ally in allies do
        if not flyingDet and ally.float then continue end

        local dif = (allies.track_pos - pos)

        if (dif.X^2 + dif.Z^2) <= rangeSquared then
            count += 1
            validIndicies[count] = ID
        end
    end
    return #validIndicies ~= 0 and validIndicies
end
local function GetAllAllies(flyingDet: boolean?)
    local indicies = {}
    local count = 0

    for ID, ally in allies do
        if not flyingDet and ally.float then continue end

        count += 1
        indicies[count] = ID
    end
    return indicies
end

local function UpdateAllySpeed(ally, value: number, increment: boolean?)
    if increment then value += ally.speed end
    ally.speed = value

    if not ally.display_speed then
        local buf = buffer.create(8)
        buffer.writeu32(buf, 0, ally.ID)
        buffer.writef32(buf, 4, value)
        remotes.UpdateAllySpeed:FireAllClients(buf)
    end
end

local function OverwriteAllySpeed(ally, value: number?, increment: boolean?)
    local buf = buffer.create(8)
    buffer.writeu32(buf, 0, ally.ID)

    if value then
        if increment then
            local displaySpeed = ally.display_speed
            if displaySpeed then
                value += displaySpeed
            end
        end

        ally.display_speed = value
        buffer.writef32(buf, 4, value)
    else
        ally.display_speed = nil
        buffer.writef32(buf, 4, ally.speed)
    end

    remotes.UpdateAllySpeed:FireAllClients(buf)
end

return {
    GetAllyCount = GetAllyCount,
    GetRandomAlly = GetRandomAlly,
    GetAlliesInRange = GetAlliesInRange,
    GetAlliesInRange2D = GetAlliesInRange2D,
    GetAllAllies = GetAllAllies,
    UpdateAllySpeed = UpdateAllySpeed,
    OverwriteAllySpeed = OverwriteAllySpeed,

    DamageAlly = DamageAlly,
    AddAlly = AddAlly,
    RemoveAlly = RemoveAlly,

    allies = allies
}