-- Collective bot budget ratelimiting.
-- Called via exec: when @pmFromFile matched a bot UA.
-- Uses user.* persistent collection (scoped per domain via setuid).
--
-- The bot table is injected at build time by do_lua_substitutions()
-- from bot_budget_bots.txt. Do not edit it manually.

-- example:
-- local BOTS = {
   -- { key = "gptbot", pattern = "GPTBot" },

-- BOT_TABLE_START
local BOTS = {
    { key = "gptbot", pattern = "GPTBot" },
    { key = "claudebot", pattern = "ClaudeBot" },
    { key = "amazonbot", pattern = "Amazonbot" },
    { key = "python_requests", pattern = "python-requests" },
    { key = "meta_externalagent", pattern = "meta-externalagent" },
    { key = "facebookcatalog", pattern = "facebookcatalog" },
    { key = "facebookexternalhit", pattern = "facebookexternalhit" },
    { key = "meta_webindexer", pattern = "meta-webindexer" }
}
-- BOT_TABLE_END

local PER_BOT_LIMIT = 20
local TOTAL_LIMIT = 50
local DRAIN_PER_MINUTE = 10
-- Drain rates (units per second, used by leaky bucket)
local PER_BOT_DRAIN = PER_BOT_LIMIT / (DRAIN_PER_MINUTE * 60)
local TOTAL_DRAIN = #BOTS * PER_BOT_DRAIN

local function get_num(key)
    local val = m.getvar(key)
    if not val or val == "" then return 0 end
    return tonumber(val) or 0
end

-- Leaky bucket: decay counter based on elapsed time, then add 1.
-- drain_rate is units per second.
local function increment_leaky(prefix, now, drain_rate)
    local count_key = "user." .. prefix .. "_count"
    local ts_key = "user." .. prefix .. "_ts"
    local count = get_num(count_key)
    local ts = get_num(ts_key)

    if ts > 0 then
        local elapsed = now - ts
        -- Drain the bucket based on time since last request
        count = math.max(0, count - (elapsed * drain_rate))
    end

    count = count + 1
    m.setvar(count_key, tostring(count))
    m.setvar(ts_key, tostring(now))
    return count
end

function main()
    local ua = m.getvar("tx.bot_ua")
    if not ua then return nil end

    local now = os.time()

    -- Identify which bot matched
    local bot_key = nil
    for _, bot in ipairs(BOTS) do
        if ua:find(bot.pattern, 1, true) then
            bot_key = bot.key
            break
        end
    end
    if not bot_key then return nil end

    -- Increment per-bot counter (leaky bucket)
    local bot_count = increment_leaky(
        "botrl_" .. bot_key, now, PER_BOT_DRAIN)

    -- Per-bot exceeded: block, don't increment total
    if bot_count > PER_BOT_LIMIT then
        m.setvar("tx.bot_block", "1")
        m.setvar("tx.bot_block_reason",
            "per-bot limit: " .. bot_key)
        return nil
    end

    -- Per-bot OK: increment total
    local total = increment_leaky(
        "botrl_total", now, TOTAL_DRAIN)

    if total > TOTAL_LIMIT then
        m.setvar("tx.bot_block", "1")
        m.setvar("tx.bot_block_reason",
            "shared budget: " .. bot_key)
    end

    return nil
end