local rand_bytes = require"kong.tools.utils".get_rand_bytes
local to_hex = require"resty.string".to_hex
local random = math.random
local kong_meta = require "kong.meta"

local TraceHandler = {PRIORITY = 1096, VERSION = kong_meta.version}

function TraceHandler:rewrite(conf)
    local b3 = kong.request.get_header("b3")
    local should_gen_trace = true
    if b3 then
        for trace_id, span_id, should_sample in string.gmatch(b3,
                                                              "(%w+)%-(%w+)-(%w+)") do
            if #trace_id == 32 and #span_id == 16 and #should_sample == 1 then
                kong.ctx.shared.trace_id = trace_id
                kong.ctx.shared.span_id = span_id
                if should_sample == "1" then
                    kong.ctx.shared.should_sample = true
                else
                    kong.ctx.shared.should_sample = false
                end
                should_gen_trace = false
            end
        end
    end
    if should_gen_trace then
        local utc = string.gsub(ngx.utctime(), "[ |%-|:]", "")
        kong.ctx.shared.trace_id = utc .. to_hex(rand_bytes(9))
        kong.ctx.shared.span_id = to_hex(rand_bytes(8))
        kong.ctx.shared.should_sample = random() < conf.sample_ratio
    end
end

function TraceHandler:access(conf)
    local b3
    if kong.ctx.shared.should_sample then
        b3 = string.format("%s-%s-%d", kong.ctx.shared.trace_id,
                           kong.ctx.shared.span_id, 1)
    else
        b3 = string.format("%s-%s-%d", kong.ctx.shared.trace_id,
                           kong.ctx.shared.span_id, 0)
    end
    kong.service.request.set_header("b3", b3)
end

-- This plugin only mutates response headers. Do not move this to the response
-- phase: Kong enables buffered proxying for every plugin that implements that
-- phase, which prevents SSE and other streaming responses from being flushed.
function TraceHandler:header_filter(conf)
    kong.response.set_header("X-Debug-Trace-Id", kong.ctx.shared.trace_id)
    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")
end

return TraceHandler