local kong_meta = require "kong.meta"
local redis = require "resty.redis"
local timestamp = require "kong.tools.timestamp"
local json = require "cjson"
json.encode_empty_table_as_object(false)

local kong = kong
local ngx = ngx
local pairs = pairs
local null = ngx.null
local fmt = string.format
local max = math.max
local time = ngx.time
local floor = math.floor
local pairs = pairs
local error = error
local tostring = tostring
local timer_at = ngx.timer.at

local EXPIRATION = {
    second = 1,
    minute = 60,
    hour = 3600,
    day = 86400,
    month = 2592000,
    year = 31536000
}

local RATELIMIT_LIMIT = "RateLimit-Limit"
local RATELIMIT_REMAINING = "RateLimit-Remaining"
local RATELIMIT_RESET = "X-RateLimit-Reset"
local RETRY_AFTER = "Retry-After"

local X_RATELIMIT_LIMIT = {
    second = "X-RateLimit-Limit-Second",
    minute = "X-RateLimit-Limit-Minute",
    hour = "X-RateLimit-Limit-Hour",
    day = "X-RateLimit-Limit-Day",
    month = "X-RateLimit-Limit-Month",
    year = "X-RateLimit-Limit-Year"
}

local X_RATELIMIT_REMAINING = {
    second = "X-RateLimit-Remaining-Second",
    minute = "X-RateLimit-Remaining-Minute",
    hour = "X-RateLimit-Remaining-Hour",
    day = "X-RateLimit-Remaining-Day",
    month = "X-RateLimit-Remaining-Month",
    year = "X-RateLimit-Remaining-Year"
}

local function gen_error_response(code, message, error_format)
    if error_format == "openapi" then
        return json.encode({code = code, message = message, details = {}})
    elseif error_format == "dataapi" then
        return json.encode({
            type = "",
            data = {
                columns = {},
                rows = {},
                result = {
                    start_ms = 0,
                    end_ms = 0,
                    latency = "",
                    row_count = 0,
                    row_affect = 0,
                    limit = 0,
                    code = code,
                    message = message
                }
            }
        })
    elseif error_format == "inner" then
        return json.encode({code = code, message = message})
    else
        return ""
    end
end

local function gen_rate_limit_message(usage, stop)
    -- Generate dynamic rate limit error message based on actual limits
    if not usage or not stop then
        return "Rate limit exceeded. Please try again later."
    end
    
    local limit = usage[stop] and usage[stop].limit
    if not limit then
        return "Rate limit exceeded. Please try again later."
    end
    
    local period_name = stop
    local request_text = limit == 1 and "request" or "requests"
    
    return fmt("Rate limit exceeded. The limit is %d %s per %s. Please try again later.", 
               limit, request_text, period_name)
end

local function is_present(str)
    return str and str ~= "" and str ~= null
end

local get_local_key = function(conf, identifier, period, period_date)
    return fmt("ratelimit:%s:%s:%s", identifier, period_date, period)
end

local sock_opts = {}
local function get_redis_connection(conf)
    local red = redis:new()
    red:set_timeout(conf.redis_timeout)

    sock_opts.ssl = conf.redis_ssl
    sock_opts.ssl_verify = conf.redis_ssl_verify

    -- use a special pool name only if redis_database is set to non-zero
    -- otherwise use the default pool name host:port
    if conf.redis_database ~= 0 then
        sock_opts.pool = fmt("%s:%d;%d", conf.redis_host, conf.redis_port,
                             conf.redis_database)
    end

    local ok, err = red:connect(conf.redis_host, conf.redis_port, sock_opts)
    if not ok then
        kong.log.err("failed to connect to Redis: ", err)
        return nil, err
    end

    local times, err = red:get_reused_times()
    if err then
        kong.log.err("failed to get connect reused times: ", err)
        return nil, err
    end

    if times == 0 then
        if is_present(conf.redis_password) then
            local ok, err
            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
                kong.log.err("failed to auth Redis: ", err)
                return nil, err
            end
        end

        if conf.redis_database ~= 0 then
            -- Only call select first time, since we know the connection is shared
            -- between instances that use the same redis database

            local ok, err = red:select(conf.redis_database)
            if not ok then
                kong.log.err("failed to change Redis database: ", err)
                return nil, err
            end
        end
    end

    return red
end

local function increment(premature, conf, limits, identifier, current_timestamp,
                         value)
    if premature then
        return
    end
    local red, err = get_redis_connection(conf)
    if not red then
        return nil, err
    end

    local keys = {}
    local expiration = {}
    local idx = 0
    local periods = timestamp.get_timestamps(current_timestamp)
    for period, period_date in pairs(periods) do
        if limits[period] then
            local cache_key = get_local_key(conf, identifier, period,
                                            period_date)
            local exists, err = red:exists(cache_key)
            if err then
                kong.log.err("failed to query Redis: ", err)
                return nil, err
            end

            idx = idx + 1
            keys[idx] = cache_key
            if not exists or exists == 0 then
                expiration[idx] = EXPIRATION[period]
            end

            red:init_pipeline()
            for i = 1, idx do
                red:incrby(keys[i], value)
                if expiration[i] then
                    red:expire(keys[i], expiration[i])
                end
            end
        end
    end

    local _, err = red:commit_pipeline()
    if err then
        kong.log.err("failed to commit increment pipeline in Redis: ", err)
        return nil, err
    end

    local ok, err = red:set_keepalive(10000, 1000)
    if not ok then
        kong.log.err("failed to set Redis keepalive: ", err)
        return nil, err
    end

    return true
end

local function policies_usage(conf, identifier, period, current_timestamp)
    local red, err = get_redis_connection(conf)
    if not red then
        return nil, err
    end

    local periods = timestamp.get_timestamps(current_timestamp)
    local cache_key = get_local_key(conf, identifier, period, periods[period])

    local current_metric, err = red:get(cache_key)
    if err then
        return nil, err
    end

    if current_metric == null then
        current_metric = nil
    end

    local ok, err = red:set_keepalive(10000, 100)
    if not ok then
        kong.log.err("failed to set Redis keepalive: ", err)
    end

    return current_metric or 0
end

local function get_identifier(conf)
    local identifier
    if conf.header_name then
        identifier = kong.request.get_header(conf.header_name)
    end
    return identifier or kong.client.get_forwarded_ip()
end

local function get_usage(conf, identifier, current_timestamp, limits)
    local usage = {}
    local stop

    for period, limit in pairs(limits) do
        local current_usage, err = policies_usage(conf, identifier, period,
                                                  current_timestamp)
        if err then
            return nil, nil, err
        end

        -- What is the current usage for the configured limit name?
        local remaining = limit - current_usage

        -- Recording usage
        usage[period] = {limit = limit, remaining = remaining}

        if remaining <= 0 then
            stop = period
        end
    end

    return usage, stop
end

local RateLimit = {PRIORITY = 1095, VERSION = kong_meta.version}

function RateLimit:access(conf)
    local current_timestamp = time() * 1000

    -- Consumer is identified by ip address or authenticated_credential id
    local identifier = get_identifier(conf)
    local fault_tolerant = conf.fault_tolerant

    -- Load current metric for configured period
    local limits = {
        second = conf.second,
        minute = conf.minute,
        hour = conf.hour,
        day = conf.day,
        month = conf.month,
        year = conf.year
    }
    -- if conf.error_format == "dataapi" and kong.ctx.shared.auth_content and
    if kong.ctx.shared.auth_content and
        kong.ctx.shared.auth_content["rate_limit_rpm"] and
        kong.ctx.shared.auth_content["rate_limit_rpm"] > 0 then
        limits = {minute = kong.ctx.shared.auth_content["rate_limit_rpm"]}
    end

    local usage, stop, err = get_usage(conf, identifier, current_timestamp,
                                       limits)
    if err then
        if not fault_tolerant then
            return error(err)
        end

        kong.log.err("failed to get usage: ", tostring(err))
    end

    if usage then
        -- Adding headers
        local reset
        local headers
        if not conf.hide_client_headers then
            headers = {}
            local timestamps
            local limit
            local window
            local remaining
            for k, v in pairs(usage) do
                local current_limit = v.limit
                local current_window = EXPIRATION[k]
                local current_remaining = v.remaining
                if stop == nil or stop == k then
                    current_remaining = current_remaining - 1
                end
                current_remaining = max(0, current_remaining)

                if not limit or (current_remaining < remaining) or
                    (current_remaining == remaining and current_window > window) then
                    limit = current_limit
                    window = current_window
                    remaining = current_remaining

                    if not timestamps then
                        timestamps = timestamp.get_timestamps(current_timestamp)
                    end

                    reset = max(1, window -
                                    floor(
                                        (current_timestamp - timestamps[k]) /
                                            1000))
                end

                headers[X_RATELIMIT_LIMIT[k]] = current_limit
                headers[X_RATELIMIT_REMAINING[k]] = current_remaining
            end

            headers[RATELIMIT_RESET] = reset
        end

        -- If limit is exceeded, terminate the request
        if stop then
            headers = headers or {}
            kong.response.set_header("X-Debug-Trace-Id",
                                     kong.ctx.shared.trace_id or "-")
            kong.response.set_header("Access-Control-Allow-Headers",
                                     "accept, content-type, authorization")
            kong.response.set_header("Access-Control-Allow-Methods",
                                     "POST, GET, OPTIONS, PUT, DELETE, PATCH")
            kong.response.set_header("Access-Control-Allow-Origin", "*")
            kong.response.add_header("Access-Control-Expose-Headers",
                                     "x-debug-trace-id,x-kong-proxy-latency,x-kong-upstream-latency,x-kong-response-latency,x-pingcap-requestid,x-ratelimit-remaining-minute,x-ratelimit-limit-minute,x-ratelimit-reset")
            headers["Content-Type"] = "application/json"
            
            -- Use dynamic error message only if explicitly configured
            local error_message
            if conf.use_dynamic_error_message then
                error_message = gen_rate_limit_message(usage, stop)
            else
                -- Keep backward compatibility with original hardcoded message
                error_message = "API key rate limit exceeded. The limit can be increased up to 1000 requests per minute per API key in TiDB Cloud console. For an increase in quota beyond 1000 rpm, please contact us: https://support.pingcap.com/hc/en-us/requests/new?ticket_form_id=7800003722519"
            end
            
            local resp = gen_error_response(49900007, error_message, conf.error_format)
            return kong.response.exit(429, resp, headers)
        end

        if headers then
            kong.response.set_headers(headers)
        end
    end

    local ok, err = timer_at(0, increment, conf, limits, identifier,
                             current_timestamp, 1)
    if not ok then
        kong.log.err("failed to create timer: ", err)
    end
end

return RateLimit