-- ============================================================
-- Macro AI: Lock villagers only (OFF-compatible)
-- ============================================================

-- ✅ global kill switch (always enable when this function is called)
__VILLAGER_LOCK_ENABLED = true

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

    local lockedSquads = {}

    ----------------------------------------------------------------
    -- Helper: lock a villager squad once (UNCHANGED PATTERN)
    ----------------------------------------------------------------
    local function LockSquadIfNeeded(sq)
        if not __VILLAGER_LOCK_ENABLED then return end
        if not sq then return end

        local squadID = Squad_GetID(sq)
        if not squadID or squadID < 0 then return end
        if lockedSquads[squadID] then return end
        if Squad_IsValid and not Squad_IsValid(squadID) then return end

        local okBP, isVill = pcall(Squad_IsOfType, sq, "scar_villager")
        if not (okBP and isVill) then return end

        AI_LockSquad(player, sq)
        lockedSquads[squadID] = true
    end

    ----------------------------------------------------------------
    -- OnUnitSpawn: lock newly spawned villagers ONLY
    -- ✅ Name matches OFF: __VillagerLock_OnSpawn
    ----------------------------------------------------------------
    function __VillagerLock_OnSpawn(event)
        if not __VILLAGER_LOCK_ENABLED then return end
        pcall(function()
            if not event or not event.entity then return end
            local squad = Entity_GetSquad(event.entity)
            if not squad then return end
            LockSquadIfNeeded(squad)
        end)
    end

    Rule_AddPlayerEvent(__VillagerLock_OnSpawn, player, GE_EntitySpawn)

    ----------------------------------------------------------------
    -- Initial lock on existing villagers (UNCHANGED PATTERN)
    ----------------------------------------------------------------
    local function _InitialLockExistingVillagers()
        if not __VILLAGER_LOCK_ENABLED then return end

        local okEg, eg = pcall(Player_GetAllEntities, player)
        if not (okEg and eg) then return end

        local total = 0
        local okCS, cntSpawn = pcall(EGroup_CountSpawned, eg)
        if okCS and type(cntSpawn) == "number" and cntSpawn > 0 then
            total = cntSpawn
        else
            local okC, cnt = pcall(EGroup_Count, eg)
            if okC and type(cnt) == "number" then total = cnt end
        end
        if total <= 0 then return end

        local function getEntityAt(i)
            if okCS and cntSpawn and cntSpawn > 0 then
                local okGetS, es = pcall(EGroup_GetSpawnedEntityAt, eg, i)
                if okGetS and es then return es end
            end
            local okGet, e2 = pcall(EGroup_GetEntityAt, eg, i)
            if okGet and e2 then return e2 end
            return nil
        end

        for i = 1, total do
            local e = getEntityAt(i)
            if e then
                local okSq, sq = pcall(Entity_GetSquad, e)
                if okSq and sq then
                    LockSquadIfNeeded(sq)
                end
            end
        end
    end

    -- Run once immediately
    pcall(_InitialLockExistingVillagers)

    -- Reset tick counter cleanly on each enable
    __EconOnly_VillagerInit_Ticks = 0

    -- Short early re-run window
    function __EconOnly_VillagerInit_Tick()
        if not __VILLAGER_LOCK_ENABLED then
            if Rule_Remove then pcall(Rule_Remove, __EconOnly_VillagerInit_Tick) end
            return
        end

        pcall(_InitialLockExistingVillagers)
        __EconOnly_VillagerInit_Ticks = (__EconOnly_VillagerInit_Ticks or 0) + 1
        if __EconOnly_VillagerInit_Ticks >= 20 then
            if Rule_Remove then pcall(Rule_Remove, __EconOnly_VillagerInit_Tick) end
        end
    end

    if Rule_AddInterval then
        Rule_AddInterval(__EconOnly_VillagerInit_Tick, 0.25)
    else
        Rule_Add(__EconOnly_VillagerInit_Tick)
    end
end

Rule_AddOneShot(InitAI_EconOnly, 1.5)