-- SAFE Combat Lock Listener (FINAL)
__ARMY_LOCK_ENABLED = true
if __ArmyLock_Installed then return end

local combatTypes = {
    -- Base military type
    ["scar_military"] = true,
    -- Infantry
    ["scar_spearman"] = true,
    ["scar_manatarms"] = true,
    ["scar_archer"] = true,
    ["scar_crossbowman"] = true,
    ["scar_repeater_crossbow"] = true,
    ["scar_handcannon"] = true,
    ["scar_landsknecht"] = true,
    ["scar_grenadier"] = true,
    -- Cavalry
    ["scar_horseman"] = true,
    ["scar_knight"] = true,
    ["scar_horsearcher"] = true,
    ["scar_camel_archer"] = true,
    ["scar_camel_rider"] = true,
    ["scar_firelancer"] = true,
    -- Unique units
    ["scar_Karash"] = true,
    ["scar_mehter"] = true,
    -- Siege
    ["scar_ram"] = true,
    ["scar_siege_tower"] = true,
    ["scar_mangonel"] = true,
    ["scar_springald"] = true,
    ["scar_ribauldequin"] = true,
    ["scar_cannon"] = true,
    ["scar_bombard"] = true,
    ["scar_culverin"] = true,
    ["scar_trebuchet"] = true,
}

__ArmyLock_LockedSquads = __ArmyLock_LockedSquads or {}

local player = Game_GetLocalPlayer()
if not player then return end

function __ArmyLock_OnSpawn(event)
    if not __ARMY_LOCK_ENABLED then return end
    
    pcall(function()
        local entity = event and event.entity
        if not entity then return end
        
        -- Skip buildings
        if type(Entity_IsBuilding) == "function" then
            local okB, isBuilding = pcall(Entity_IsBuilding, entity)
            if okB and isBuilding then return end
        end
        
        local okSquad, squad = pcall(Entity_GetSquad, entity)
        if not okSquad or not squad then return end
        
        -- Skip dead/dying squads
        if type(Squad_IsAlive) == "function" then
            local okAlive, alive = pcall(Squad_IsAlive, squad)
            if not okAlive or not alive then return end
        end
        
        -- REMOVED Squad_IsValid check - it fails for newly spawned squads
        
        local okID, squadID = pcall(Squad_GetID, squad)
        if not okID or not squadID or squadID < 0 then return end
        
        if __ArmyLock_LockedSquads[squadID] then return end
        
        -- Check if combat type
        local isCombat = false
        for t,_ in pairs(combatTypes) do
            if Squad_IsOfType then
                local okType, isType = pcall(Squad_IsOfType, squad, t)
                if okType and isType then
                    isCombat = true
                    break
                end
            end
        end
        
        if not isCombat then return end
        
        -- Lock the squad
        local okLock = pcall(AI_LockSquad, player, squad)
        if okLock then
            __ArmyLock_LockedSquads[squadID] = true
        end
    end)
end

Rule_AddPlayerEvent(__ArmyLock_OnSpawn, player, GE_EntitySpawn)
__ArmyLock_Installed = true