#!/usr/bin/env tarantool
-- Microbenchmark harness for protoc-gen-tarantool.
--
-- Measures encode + decode throughput and allocation rate across 5 payload
-- sizes (~10 B, ~100 B, ~1 KB, ~10 KB, ~100 KB) for both codegen modes
-- (full inline / runtime descriptor). Emits a JSON document on stdout that
-- can be compared against `bench/baseline.json`.
--
-- Usage:
--   tarantool bench/bench.lua                  -- run, print JSON
--   tarantool bench/bench.lua --baseline       -- overwrite baseline.json
--   tarantool bench/bench.lua --compare        -- compare vs baseline.json
--                                                 exit nonzero if any
--                                                 throughput regresses >5%

package.path = './runtime/?.lua;./runtime/?/init.lua;'
    .. './examples/expected/?.lua;./examples/expected/?/init.lua;'
    .. package.path

local clock = require('clock')
local json  = require('json')
local fio   = require('fio')

local MODES = {'full', 'runtime'}
local SIZES = {
    {label = '10B',   target = 10},
    {label = '100B',  target = 100},
    {label = '1KB',   target = 1024},
    {label = '10KB',  target = 10240},
    {label = '100KB', target = 102400},
}

-- Build a `Person` payload whose encoded size is close to `target` bytes.
--
-- Strategy: pick one knob per decade so each size still exercises the
-- full encoder (varints, packed repeated, length-delimited strings,
-- nested messages) — not just one giant byte-string.
local function build_payload(target)
    if target <= 10 then
        -- name(6) + age(1) ⇒ 10 bytes encoded.
        return {name = 'bigbes', age = 42}
    end
    if target <= 100 then
        -- name (~target-10 bytes string) gives a tight fit (~94 B).
        return {
            name  = string.rep('a', target - 10),
            age   = 42,
        }
    end
    -- For >=1 KB: scale `emails` (length-delimited strings) and add nested
    -- + packed repeated fields so the shape stays representative.
    local per_email   = 36  -- tag(1) + len(1) + 32 bytes content + slack
    local fixed_bytes = 80  -- name + age + address + lucky_numbers + overhead
    local n_emails    = math.max(1, math.floor((target - fixed_bytes) / per_email))
    local p = {
        name = 'bigbes',
        age  = 42,
        address = {street = '1 Main St', city = 'Springfield', zip = 12345},
        lucky_numbers = {7, 13, 21, 42, 99},
        emails = {},
    }
    for i = 1, n_emails do
        p.emails[i] = string.rep('e', 28) .. string.format('%04d', i)
    end
    return p
end

-- Pick iteration count adaptively: smaller messages need more iters to
-- amortize loop + clock overhead; larger messages need fewer to keep
-- wall time bounded.
local function iter_count(size_bytes)
    if size_bytes <  100   then return 200000 end
    if size_bytes <  2000  then return  50000 end
    if size_bytes <  20000 then return   5000 end
    return 500
end

-- Median + min/max from a small sample. Median rejects single-trace
-- compilation outliers; min is closer to steady-state JIT performance.
local function summarize(samples)
    table.sort(samples)
    local n = #samples
    local median = samples[math.floor((n + 1) / 2)]
    return {
        median = median,
        min    = samples[1],
        max    = samples[n],
    }
end

local function time_loop(fn, n)
    local t0 = clock.monotonic64()
    for _ = 1, n do fn() end
    local t1 = clock.monotonic64()
    return tonumber(t1 - t0) / 1e9  -- seconds
end

local function bench_throughput(fn, n, runs)
    -- Warmup: let the JIT compile.
    for _ = 1, math.min(n, 1000) do fn() end
    local times = {}
    for r = 1, runs do
        collectgarbage('collect')
        times[r] = time_loop(fn, n)
    end
    local s = summarize(times)
    return {
        ns_per_op   = s.median / n * 1e9,
        msgs_per_s  = n / s.median,
        runs        = runs,
        iters       = n,
        time_min_s  = s.min,
        time_med_s  = s.median,
        time_max_s  = s.max,
    }
end

-- Allocation per op. Stop GC, run a small batch, measure delta in KB.
-- Restart GC immediately so the next bench isn't polluted.
--
-- Iteration count is capped so peak retained memory stays under ~64 MB
-- — for 100KB messages 1000 iters would hold 500MB live and trigger OS
-- swap pressure that skews adjacent throughput readings.
local function bench_alloc(fn, expected_bytes)
    local budget   = 64 * 1024 * 1024
    local per_iter = math.max(1, expected_bytes) * 2
    local n        = math.max(100, math.min(2000, math.floor(budget / per_iter)))
    -- prime: ensure any one-shot allocations (descriptor lookups, jit
    -- traces) already happened.
    for _ = 1, 100 do fn() end
    collectgarbage('collect')
    collectgarbage('stop')
    local before = collectgarbage('count')
    for _ = 1, n do fn() end
    local after = collectgarbage('count')
    collectgarbage('restart')
    collectgarbage('collect')
    return {
        kb_per_op    = (after - before) / n,
        bytes_per_op = (after - before) * 1024 / n,
        iters        = n,
    }
end

local function bench_one(mode, size)
    local hello_pb = require(mode .. '.hello.hello_pb')
    local encode   = hello_pb.Person_encode
    local decode   = hello_pb.Person_decode

    local payload = build_payload(size.target)
    local bytes   = encode(payload)
    local n       = iter_count(#bytes)
    local runs    = 5

    -- Re-decode once so warmup hot path matches.
    local _ = decode(bytes)

    local enc_throughput = bench_throughput(function() encode(payload) end, n, runs)
    local enc_alloc      = bench_alloc(function() encode(payload) end, #bytes)
    enc_throughput.mb_per_s = #bytes * enc_throughput.msgs_per_s / 1e6
    enc_throughput.alloc_kb_per_op    = enc_alloc.kb_per_op
    enc_throughput.alloc_bytes_per_op = enc_alloc.bytes_per_op

    local dec_throughput = bench_throughput(function() decode(bytes) end, n, runs)
    local dec_alloc      = bench_alloc(function() decode(bytes) end, #bytes)
    dec_throughput.mb_per_s = #bytes * dec_throughput.msgs_per_s / 1e6
    dec_throughput.alloc_kb_per_op    = dec_alloc.kb_per_op
    dec_throughput.alloc_bytes_per_op = dec_alloc.bytes_per_op

    return {
        mode       = mode,
        size_label = size.label,
        size_bytes = #bytes,
        encode     = enc_throughput,
        decode     = dec_throughput,
    }
end

local function run_all()
    local results = {}
    for _, mode in ipairs(MODES) do
        for _, size in ipairs(SIZES) do
            io.stderr:write(string.format('  bench %s/%s ... ', mode, size.label))
            io.stderr:flush()
            local r = bench_one(mode, size)
            io.stderr:write(string.format(
                'enc %.0f msgs/s (%.1f MB/s)  dec %.0f msgs/s (%.1f MB/s)\n',
                r.encode.msgs_per_s, r.encode.mb_per_s,
                r.decode.msgs_per_s, r.decode.mb_per_s))
            results[#results + 1] = r
        end
    end
    return {
        tarantool       = _TARANTOOL,
        jit             = jit and jit.version or nil,
        schema_message  = 'hello.Person',
        results         = results,
    }
end

-- Render JSON deterministically: arrays preserve order, but Lua tables
-- iterate in hash order. We control key emission for each level.
local function render_metric(t)
    return string.format(
        '{"ns_per_op": %.1f, "msgs_per_s": %.0f, "mb_per_s": %.3f, '
        .. '"alloc_kb_per_op": %.3f, "alloc_bytes_per_op": %.1f, '
        .. '"iters": %d, "runs": %d, "time_med_s": %.6f, '
        .. '"time_min_s": %.6f, "time_max_s": %.6f}',
        t.ns_per_op, t.msgs_per_s, t.mb_per_s,
        t.alloc_kb_per_op, t.alloc_bytes_per_op,
        t.iters, t.runs, t.time_med_s, t.time_min_s, t.time_max_s)
end

local function render(doc)
    local lines = {}
    lines[#lines + 1] = '{'
    lines[#lines + 1] = string.format('  "tarantool": %s,', json.encode(doc.tarantool))
    lines[#lines + 1] = string.format('  "jit": %s,', json.encode(doc.jit or json.NULL))
    lines[#lines + 1] = string.format('  "schema_message": %s,', json.encode(doc.schema_message))
    lines[#lines + 1] = '  "results": ['
    for i, r in ipairs(doc.results) do
        local sep = (i == #doc.results) and '' or ','
        lines[#lines + 1] = '    {'
        lines[#lines + 1] = string.format('      "mode": %s,', json.encode(r.mode))
        lines[#lines + 1] = string.format('      "size_label": %s,', json.encode(r.size_label))
        lines[#lines + 1] = string.format('      "size_bytes": %d,', r.size_bytes)
        lines[#lines + 1] = string.format('      "encode": %s,', render_metric(r.encode))
        lines[#lines + 1] = string.format('      "decode": %s', render_metric(r.decode))
        lines[#lines + 1] = '    }' .. sep
    end
    lines[#lines + 1] = '  ]'
    lines[#lines + 1] = '}'
    return table.concat(lines, '\n') .. '\n'
end

-- Hardware-portable baseline: throughput (msgs/s, MB/s) varies with CPU
-- load and is unsuitable for committed baselines. Allocation per op is
-- reproducible to within ~10 bytes regardless of machine — it counts
-- bytes, not time — so that's all we commit. Throughput is in --print
-- output for human inspection only.
local function reduce_for_baseline(doc)
    local by_key = {}
    for _, r in ipairs(doc.results) do
        by_key[r.mode .. '/' .. r.size_label] = r
    end
    local out = {}
    for _, size in ipairs(SIZES) do
        local full    = by_key['full/'    .. size.label]
        local runtime = by_key['runtime/' .. size.label]
        out[#out + 1] = {
            size_label = size.label,
            size_bytes = full.size_bytes,
            encode = {
                alloc_kb_per_op_full    = full.encode.alloc_kb_per_op,
                alloc_kb_per_op_runtime = runtime.encode.alloc_kb_per_op,
            },
            decode = {
                alloc_kb_per_op_full    = full.decode.alloc_kb_per_op,
                alloc_kb_per_op_runtime = runtime.decode.alloc_kb_per_op,
            },
        }
    end
    return {schema_message = doc.schema_message, results = out}
end

local function render_baseline(reduced)
    local lines = {'{'}
    lines[#lines + 1] = string.format('  "schema_message": %s,', json.encode(reduced.schema_message))
    lines[#lines + 1] = '  "results": ['
    for i, r in ipairs(reduced.results) do
        local sep = (i == #reduced.results) and '' or ','
        lines[#lines + 1] = '    {'
        lines[#lines + 1] = string.format('      "size_label": %s,', json.encode(r.size_label))
        lines[#lines + 1] = string.format('      "size_bytes": %d,', r.size_bytes)
        lines[#lines + 1] = string.format(
            '      "encode": {"alloc_kb_per_op_full": %.3f, '
            .. '"alloc_kb_per_op_runtime": %.3f},',
            r.encode.alloc_kb_per_op_full,
            r.encode.alloc_kb_per_op_runtime)
        lines[#lines + 1] = string.format(
            '      "decode": {"alloc_kb_per_op_full": %.3f, '
            .. '"alloc_kb_per_op_runtime": %.3f}',
            r.decode.alloc_kb_per_op_full,
            r.decode.alloc_kb_per_op_runtime)
        lines[#lines + 1] = '    }' .. sep
    end
    lines[#lines + 1] = '  ]'
    lines[#lines + 1] = '}'
    return table.concat(lines, '\n') .. '\n'
end

-- Compare two reduced baselines, return list of regressions exceeding
-- `tolerance` (fraction, e.g. 0.05 = 5%).
--
-- Allocation per op is the regression gate. It's hardware-independent
-- (counts bytes, not time), reproducible to within ~10 bytes per op,
-- and a direct measure of encoder/decoder efficiency. Throughput
-- ratios swing 30%+ run-to-run on a busy laptop — useless as a gate.
local function compare(current, baseline, tolerance)
    local function index(b)
        local m = {}
        for _, r in ipairs(b.results) do m[r.size_label] = r end
        return m
    end
    local cur = index(current)
    local base = index(baseline)
    local regressions = {}
    for _, size in ipairs(SIZES) do
        local c = cur[size.label]
        local b = base[size.label]
        if not c or not b then goto continue end
        for _, op in ipairs({'encode', 'decode'}) do
            for _, key in ipairs({'alloc_kb_per_op_full', 'alloc_kb_per_op_runtime'}) do
                local bv, cv = b[op][key], c[op][key]
                if bv > 0 and cv > bv * (1 + tolerance) then
                    regressions[#regressions + 1] = string.format(
                        '%s/%s %s: %.3f -> %.3f KB/op (+%.1f%%)',
                        size.label, op, key, bv, cv, (cv / bv - 1) * 100)
                end
            end
        end
        ::continue::
    end
    return regressions
end

local args = {...}
local mode_flag = args[1] or '--print'

io.stderr:write(string.format('tarantool-protobuf bench (%s)\n', _TARANTOOL))
local doc = run_all()

if mode_flag == '--print' then
    io.write(render(doc))
elseif mode_flag == '--baseline' then
    local out = render_baseline(reduce_for_baseline(doc))
    local path = 'bench/baseline.json'
    local f = assert(fio.open(path, {'O_WRONLY', 'O_CREAT', 'O_TRUNC'}, tonumber('644', 8)))
    f:write(out)
    f:close()
    io.stderr:write(string.format('wrote %s\n', path))
    io.write(out)
elseif mode_flag == '--compare' then
    local path = 'bench/baseline.json'
    local f = assert(fio.open(path, {'O_RDONLY'}))
    local baseline = json.decode(f:read())
    f:close()
    local current = reduce_for_baseline(doc)
    local regs = compare(current, baseline, 0.05)
    if #regs == 0 then
        io.stderr:write('no regressions >5% vs baseline\n')
        os.exit(0)
    end
    io.stderr:write(string.format('REGRESSIONS vs baseline (>5%%):\n'))
    for _, r in ipairs(regs) do io.stderr:write('  ' .. r .. '\n') end
    os.exit(1)
else
    io.stderr:write('usage: bench.lua [--print | --baseline | --compare]\n')
    os.exit(2)
end

os.exit(0)