local cjson = require "cjson.safe"
local kong_meta = require "kong.meta"
local redis = require "resty.redis"
local resty_sha256 = require "resty.sha256"
local resty_string = require "resty.string"

local kong = kong
local ngx = ngx

local ACQUIRE_SCRIPT = [[
local now = tonumber(ARGV[5])
redis.call("ZREMRANGEBYSCORE", KEYS[1], "-inf", now)
redis.call("ZREMRANGEBYSCORE", KEYS[3], "-inf", now)
if ARGV[8] == "1" then
  redis.call("ZREMRANGEBYSCORE", KEYS[2], "-inf", now)
end

local rate = tonumber(redis.call("GET", KEYS[4]) or "0")
if rate >= tonumber(ARGV[4]) then
  return {0, "rate", rate}
end

redis.call("INCR", KEYS[4])
redis.call("EXPIRE", KEYS[4], 61)

local global_active = tonumber(redis.call("ZCARD", KEYS[1]) or "0")
if global_active >= tonumber(ARGV[1]) then
  return {0, "global", global_active}
end

local user_active = tonumber(redis.call("ZCARD", KEYS[3]) or "0")
if user_active >= tonumber(ARGV[3]) then
  return {0, "user", user_active}
end

if ARGV[8] == "1" then
  local org_active = tonumber(redis.call("ZCARD", KEYS[2]) or "0")
  if org_active >= tonumber(ARGV[2]) then
    return {0, "org", org_active}
  end
end

redis.call("ZADD", KEYS[1], ARGV[6], ARGV[9])
redis.call("EXPIRE", KEYS[1], ARGV[7])
redis.call("ZADD", KEYS[3], ARGV[6], ARGV[9])
redis.call("EXPIRE", KEYS[3], ARGV[7])

if ARGV[8] == "1" then
  redis.call("ZADD", KEYS[2], ARGV[6], ARGV[9])
  redis.call("EXPIRE", KEYS[2], ARGV[7])
end

return {1, "ok", global_active + 1, user_active + 1}
]]

local RELEASE_SCRIPT = [[
for i = 1, #KEYS do
  redis.call("ZREM", KEYS[i], ARGV[1])
  if redis.call("ZCARD", KEYS[i]) == 0 then
    redis.call("DEL", KEYS[i])
  end
end
return 1
]]

local RENEW_SCRIPT = [[
local renewed = 0
for i = 1, #KEYS do
  if redis.call("ZSCORE", KEYS[i], ARGV[1]) then
    redis.call("ZADD", KEYS[i], "XX", ARGV[2], ARGV[1])
    redis.call("EXPIRE", KEYS[i], ARGV[3])
    renewed = renewed + 1
  end
end
return renewed
]]

local function is_present(value)
    return value and value ~= "" and value ~= ngx.null
end

local function hash_key(value)
    local sha256 = resty_sha256:new()
    if not sha256 then
        error("failed to initialize SHA-256")
    end
    local ok = sha256:update(tostring(value))
    if not ok then
        error("failed to update SHA-256")
    end
    return resty_string.to_hex(sha256:final())
end

local function redis_connection(conf)
    local red = redis:new()
    red:set_timeout(conf.redis_timeout)

    local options = {
        ssl = conf.redis_ssl,
        ssl_verify = conf.redis_ssl_verify
    }
    if conf.redis_database ~= 0 then
        options.pool = string.format("%s:%d;%d", conf.redis_host,
                                     conf.redis_port, conf.redis_database)
    end

    local ok, err = red:connect(conf.redis_host, conf.redis_port, options)
    if not ok then
        return nil, err
    end

    local reused, reused_err = red:get_reused_times()
    if reused_err then
        return nil, reused_err
    end

    if reused == 0 then
        if is_present(conf.redis_password) then
            if is_present(conf.redis_username) then
                ok, err = red:auth(conf.redis_username, conf.redis_password)
            else
                ok, err = red:auth(conf.redis_password)
            end
            if not ok then
                return nil, err
            end
        end

        if conf.redis_database ~= 0 then
            ok, err = red:select(conf.redis_database)
            if not ok then
                return nil, err
            end
        end
    end

    return red
end

local function keepalive(red)
    local ok, err = red:set_keepalive(10000, 1000)
    if not ok then
        ngx.log(ngx.ERR,
                "investibook SSE limiter failed to keep Redis connection alive: ",
                err)
    end
end

local function trusted_user_id(conf)
    local user_id = kong.ctx.shared.investibook_user_id
    if not is_present(user_id) and conf.identity_header_name then
        user_id = kong.request.get_header(conf.identity_header_name)
    end
    if not is_present(user_id) then
        user_id = kong.client.get_forwarded_ip()
    end
    return tostring(user_id or "unknown")
end

local function org_id(conf)
    local value = kong.ctx.shared.investibook_org_id
    if not is_present(value) and conf.org_header_name then
        -- This header is only a fairness partition. The authenticated user and
        -- global limits remain the security boundary if the header is absent
        -- or caller-controlled.
        value = kong.request.get_header(conf.org_header_name)
    end
    if is_present(value) then
        return tostring(value)
    end
    return nil
end

local function reject(status, detail, retry_after)
    return kong.response.exit(status, cjson.encode({detail = detail}), {
        ["Content-Type"] = "application/json",
        ["Retry-After"] = tostring(retry_after)
    })
end

local InvestibookSSELimit = {
    PRIORITY = 1094,
    VERSION = kong_meta.version
}

local release_lease

local function is_sse_request()
    local accept = kong.request.get_header("accept")
    if type(accept) ~= "string" then
        return false
    end
    return string.find(string.lower(accept), "text/event-stream", 1, true) ~= nil
end

local function heartbeat_lease(premature, conf, release)
    if premature or not release.active then
        return
    end

    local red, err = redis_connection(conf)
    if not red then
        ngx.log(ngx.ERR,
                "investibook SSE limiter failed to renew Redis lease: ", err)
    else
        local keys = release.keys
        local arguments = {RENEW_SCRIPT, #keys}
        for _, key in ipairs(keys) do
            table.insert(arguments, key)
        end
        table.insert(arguments, release.lease_id)
        table.insert(arguments, ngx.time() + conf.lease_ttl_seconds)
        table.insert(arguments, conf.lease_ttl_seconds * 2)

        local renewed, eval_err = red:eval(unpack(arguments))
        if eval_err then
            ngx.log(ngx.ERR,
                    "investibook SSE limiter Redis lease renewal failed: ",
                    eval_err)
        elseif tonumber(renewed) ~= #keys and release.active then
            ngx.log(ngx.WARN,
                    "investibook SSE limiter renewed only ", renewed, " of ",
                    #keys, " Redis lease indexes")
        end
        keepalive(red)
    end

    if release.active then
        local ok, timer_err = ngx.timer.at(release.heartbeat_interval,
                                           heartbeat_lease, conf, release)
        if not ok then
            ngx.log(ngx.ERR,
                    "investibook SSE limiter failed to reschedule lease heartbeat: ",
                    timer_err)
        end
    end
end

function InvestibookSSELimit:access(conf)
    if kong.request.get_method() == "OPTIONS" or not is_sse_request() then
        return
    end

    local user_hash = hash_key(trusted_user_id(conf))
    local current_org_id = org_id(conf)
    local has_org = current_org_id ~= nil
    local org_hash = has_org and hash_key(current_org_id) or "none"
    local prefix = conf.key_prefix
    local minute_bucket = math.floor(ngx.time() / 60)
    local now = ngx.time()
    local expires_at = now + conf.lease_ttl_seconds
    local request_id = kong.ctx.shared.trace_id or ngx.var.request_id or
        tostring(ngx.now()) .. ":" .. tostring(math.random())
    local lease_id = hash_key(user_hash .. ":" .. request_id)
    local keys = {
        prefix .. ":active:global",
        prefix .. ":active:org:" .. org_hash,
        prefix .. ":active:user:" .. user_hash,
        prefix .. ":opens:user:" .. user_hash .. ":" .. minute_bucket
    }

    local red, err = redis_connection(conf)
    if not red then
        kong.log.err("investibook SSE limiter failed to connect to Redis: ", err)
        if conf.fault_tolerant then
            return
        end
        return reject(503, "Streaming capacity protection is temporarily unavailable.",
                      conf.retry_after_seconds)
    end

    local result, eval_err = red:eval(ACQUIRE_SCRIPT, #keys, keys[1], keys[2],
                                      keys[3], keys[4],
                                      conf.global_concurrency,
                                      conf.org_concurrency,
                                      conf.user_concurrency,
                                      conf.connections_per_minute, now,
                                      expires_at,
                                      conf.lease_ttl_seconds * 2,
                                      has_org and 1 or 0, lease_id)
    keepalive(red)

    if not result then
        kong.log.err("investibook SSE limiter Redis acquire failed: ", eval_err)
        if conf.fault_tolerant then
            return
        end
        return reject(503, "Streaming capacity protection is temporarily unavailable.",
                      conf.retry_after_seconds)
    end

    if tonumber(result[1]) ~= 1 then
        kong.log.warn("investibook SSE connection rejected by ", result[2],
                      " limit")
        return reject(429, "Too many streaming connections. Please retry later.",
                      conf.retry_after_seconds)
    end

    local release_keys = {keys[1], keys[3]}
    if has_org then
        table.insert(release_keys, keys[2])
    end
    local release = {
        keys = release_keys,
        lease_id = lease_id,
        active = true,
        heartbeat_interval = math.max(0.1, conf.lease_ttl_seconds / 3)
    }
    kong.ctx.shared.investibook_sse_limit_release = release

    local heartbeat_ok, heartbeat_err = ngx.timer.at(
                                              release.heartbeat_interval,
                                              heartbeat_lease, conf, release)
    if not heartbeat_ok then
        release.active = false
        kong.log.err(
            "investibook SSE limiter failed to schedule lease heartbeat: ",
            heartbeat_err)
        release_lease(false, conf, release)
        kong.ctx.shared.investibook_sse_limit_release = nil
        return reject(503,
                      "Streaming capacity protection is temporarily unavailable.",
                      conf.retry_after_seconds)
    end
end

release_lease = function(premature, conf, release)
    if premature then
        return
    end

    local red, err = redis_connection(conf)
    if not red then
        ngx.log(ngx.ERR,
                "investibook SSE limiter failed to release Redis lease: ", err)
        return
    end

    local keys = release.keys
    local arguments = {RELEASE_SCRIPT, #keys}
    for _, key in ipairs(keys) do
        table.insert(arguments, key)
    end
    table.insert(arguments, release.lease_id)
    local _, eval_err = red:eval(unpack(arguments))
    if eval_err then
        ngx.log(ngx.ERR, "investibook SSE limiter Redis release failed: ",
                eval_err)
    end
    keepalive(red)
end

function InvestibookSSELimit:log(conf)
    local release = kong.ctx.shared.investibook_sse_limit_release
    if not release or not release.keys or #release.keys == 0 then
        return
    end

    release.active = false

    -- The log phase cannot use cosockets directly. Schedule Redis cleanup in
    -- a zero-delay timer, where the cosocket API is available.
    local ok, err = ngx.timer.at(0, release_lease, conf, release)
    if not ok then
        kong.log.err("investibook SSE limiter failed to schedule lease release: ",
                     err)
    end
end

return InvestibookSSELimit