-- Low-level protobuf wire format (proto3).
-- Pure Lua + LuaJIT FFI; no Tarantool-specific dependencies.
local ffi = require('ffi')
local bit = require('bit')

local M = {}

-- Wire type constants (https://protobuf.dev/programming-guides/encoding/#structure)
M.WIRE_VARINT = 0
M.WIRE_I64    = 1
M.WIRE_LEN    = 2
M.WIRE_I32    = 5

local UINT64 = ffi.typeof('uint64_t')
local INT64  = ffi.typeof('int64_t')
local UINT64_ZERO = UINT64(0)
local CONT_MASK   = UINT64(bit.bnot(0x7f))  -- 0xFFFFFFFFFFFFFF80

-- Coerce any integer-like value to uint64_t cdata.
-- Negative Lua numbers are sign-extended via int64_t (protobuf wire spec).
local function to_uint64(v)
    local t = type(v)
    if t == 'number' then
        if v < 0 then return UINT64(INT64(v)) end
        return UINT64(v)
    elseif t == 'cdata' then
        return UINT64(v)
    elseif t == 'boolean' then
        return v and UINT64(1) or UINT64_ZERO
    end
    error("cannot coerce " .. t .. " to uint64", 0)
end
M.to_uint64 = to_uint64

local function to_int64(v)
    local t = type(v)
    if t == 'number' or t == 'cdata' then return INT64(v) end
    error("cannot coerce " .. t .. " to int64", 0)
end
M.to_int64 = to_int64

-- Truncate a uint64 varint payload to proto3 int32/uint32 Lua numbers.
-- Per spec, int32/uint32/enum/sint32 fields must keep only the low 32 bits
-- of an over-range varint; int32/sint32 additionally sign-extend from bit 31.
local function varint_to_uint32(u)
    return tonumber(bit.band(u, 0xFFFFFFFF))
end
local function varint_to_int32(u)
    local n = tonumber(bit.band(u, 0xFFFFFFFF))
    if n >= 0x80000000 then n = n - 0x100000000 end
    return n
end
M.varint_to_uint32 = varint_to_uint32
M.varint_to_int32  = varint_to_int32

-- ---------------------------------------------------------------------------
-- Varint
-- ---------------------------------------------------------------------------

-- encode_varint(n) -> string
-- Accepts uint64_t/int64_t cdata, Lua number, or boolean.
--
-- Fast path: small non-negative Lua numbers (0..127) become a single
-- string.char(n) call with no cdata allocation, no `out` table, no
-- table.concat. Covers most length prefixes for short strings, many
-- enum ordinals, and most small int values in typical RPC payloads.
--
-- We deliberately do NOT extend the fast path to 2-4 byte values:
-- growing the function past the LuaJIT inline budget makes parent
-- traces stop inlining it, which costs more (~30% bench regression
-- on 1-byte-dominant workloads) than the rare multi-byte case gains.
local function encode_varint(n)
    if type(n) == 'number' and n >= 0 and n < 0x80 then
        return string.char(n)
    end
    n = to_uint64(n)
    local out = {}
    local i = 1
    while bit.band(n, CONT_MASK) ~= UINT64_ZERO do
        out[i] = string.char(tonumber(bit.bor(bit.band(n, 0x7f), 0x80)))
        n = bit.rshift(n, 7)
        i = i + 1
    end
    out[i] = string.char(tonumber(n))
    return table.concat(out)
end
M.encode_varint = encode_varint

-- decode_varint(buf, pos) -> uint64_t cdata, new_pos (1-based)
--
-- Fast path is inlined: 1-byte varints (field tags for ids 1..15 and
-- many small values) take a straight-line branch with no loop, which
-- keeps the JIT trace single-rooted across hot decode callers. The
-- multi-byte tail still loops, but it's only entered for the small
-- minority of values that don't fit in 7 bits.
local function decode_varint(buf, pos)
    local b = buf:byte(pos)
    if b == nil then error("truncated varint at offset " .. pos, 0) end
    if b < 0x80 then
        return UINT64(b), pos + 1
    end
    local result = UINT64(bit.band(b, 0x7f))
    local shift = 7
    pos = pos + 1
    while true do
        b = buf:byte(pos)
        if b == nil then error("truncated varint at offset " .. pos, 0) end
        pos = pos + 1
        result = bit.bor(result, bit.lshift(UINT64(bit.band(b, 0x7f)), shift))
        if b < 0x80 then return result, pos end
        shift = shift + 7
        if shift >= 70 then error("varint exceeds 10 bytes", 0) end
    end
end
M.decode_varint = decode_varint

-- ---------------------------------------------------------------------------
-- Tag
-- ---------------------------------------------------------------------------

local function encode_tag(field_id, wire_type)
    -- field_id < 2^29, fits in Lua double exactly.
    return encode_varint(field_id * 8 + wire_type)
end
M.encode_tag = encode_tag

-- The 1-byte varint fast path is duplicated at every hot decode call
-- site (decode_tag, decode_len, decode_int32 / int64 / uint32 / uint64 /
-- sint32 / sint64 / bool, skip_field VARINT branch) instead of being
-- factored into a helper. Reason: LuaJIT 2.1 inlines a small called
-- function into the caller's trace, so when the parent guard exits to
-- a side trace for the multi-byte case, the side trace has to return
-- from the inlined frame — and LuaJIT can't stitch that return back
-- to the parent, dropping to interpreter dispatch. Inlining the fast
-- path literally keeps the side trace inside the caller's own frame,
-- where stitching works.
local function decode_tag(buf, pos)
    local b = buf:byte(pos)
    if b == nil then error("truncated varint at offset " .. pos, 0) end
    if b < 0x80 then
        local wt = bit.band(b, 7)
        if wt >= 6 then error("illegal wire type " .. wt, 0) end
        local fn = bit.rshift(b, 3)
        if fn == 0 then error("illegal field number 0", 0) end
        return fn, wt, pos + 1
    end
    local u, npos = decode_varint(buf, pos)
    -- A tag must use the minimum number of bytes to encode its value.
    -- The trailing byte of a multi-byte varint contributes 0 high bits
    -- only when the encoding is overlong (any prior byte already covered
    -- the value).
    if buf:byte(npos - 1) == 0 then
        error("overlong tag varint at offset " .. pos, 0)
    end
    -- Bit ops on the uint64 cdata preserve 64-bit width; going through
    -- tonumber first would truncate field numbers above 2^32.
    local wt = tonumber(bit.band(u, 7))
    if wt >= 6 then error("illegal wire type " .. wt, 0) end
    local fn = tonumber(bit.rshift(u, 3))
    if fn == 0 then error("illegal field number 0", 0) end
    -- Field numbers are 29-bit per the protobuf spec.
    if fn > 0x1FFFFFFF then
        error("field number out of range: " .. fn, 0)
    end
    return fn, wt, npos
end
M.decode_tag = decode_tag

-- ---------------------------------------------------------------------------
-- ZigZag (sint32 / sint64)
-- ---------------------------------------------------------------------------

-- 32-bit zigzag stays in Lua number range.
local function zigzag_encode32(n)
    n = tonumber(n)
    if n >= 0 then return n * 2 else return -n * 2 - 1 end
end
M.zigzag_encode32 = zigzag_encode32

local function zigzag_decode32(u)
    u = tonumber(u)
    if u % 2 == 0 then return u / 2 else return -((u + 1) / 2) end
end
M.zigzag_decode32 = zigzag_decode32

-- 64-bit zigzag uses cdata.
local function zigzag_encode64(n)
    local i = to_int64(n)
    local doubled = bit.lshift(UINT64(i), 1)
    if i >= 0 then return doubled end
    return bit.bnot(doubled)
end
M.zigzag_encode64 = zigzag_encode64

local function zigzag_decode64(u)
    u = to_uint64(u)
    local half = bit.rshift(u, 1)
    if bit.band(u, 1) == UINT64_ZERO then return INT64(half) end
    return INT64(bit.bnot(half))
end
M.zigzag_decode64 = zigzag_decode64

-- ---------------------------------------------------------------------------
-- Fixed32 / Fixed64 (little-endian)
-- ---------------------------------------------------------------------------

local function encode_fixed32(v)
    -- Accepts uint32 (Lua number 0..2^32-1) or any integer cdata.
    local u
    if type(v) == 'cdata' then
        u = tonumber(bit.band(UINT64(v), 0xffffffff))
    else
        u = tonumber(v)
        if u < 0 then u = u + 0x100000000 end
    end
    return string.char(
        bit.band(u, 0xff),
        bit.band(bit.rshift(u, 8), 0xff),
        bit.band(bit.rshift(u, 16), 0xff),
        bit.band(bit.rshift(u, 24), 0xff))
end
M.encode_fixed32 = encode_fixed32

-- decode_fixed32(buf, pos) -> Lua number (0..2^32-1), new_pos
local function decode_fixed32(buf, pos)
    local b1, b2, b3, b4 = buf:byte(pos, pos + 3)
    if b4 == nil then error("truncated fixed32", 0) end
    return b1 + b2 * 0x100 + b3 * 0x10000 + b4 * 0x1000000, pos + 4
end
M.decode_fixed32 = decode_fixed32

local function encode_fixed64(v)
    local u = to_uint64(v)
    local lo = tonumber(bit.band(u, 0xffffffff))
    local hi = tonumber(bit.rshift(u, 32))
    return string.char(
        bit.band(lo, 0xff),
        bit.band(bit.rshift(lo, 8), 0xff),
        bit.band(bit.rshift(lo, 16), 0xff),
        bit.band(bit.rshift(lo, 24), 0xff),
        bit.band(hi, 0xff),
        bit.band(bit.rshift(hi, 8), 0xff),
        bit.band(bit.rshift(hi, 16), 0xff),
        bit.band(bit.rshift(hi, 24), 0xff))
end
M.encode_fixed64 = encode_fixed64

-- decode_fixed64(buf, pos) -> uint64_t cdata, new_pos
local function decode_fixed64(buf, pos)
    local b1, b2, b3, b4, b5, b6, b7, b8 = buf:byte(pos, pos + 7)
    if b8 == nil then error("truncated fixed64", 0) end
    local lo = b1 + b2 * 0x100 + b3 * 0x10000 + b4 * 0x1000000
    local hi = b5 + b6 * 0x100 + b7 * 0x10000 + b8 * 0x1000000
    return UINT64(lo) + bit.lshift(UINT64(hi), 32), pos + 8
end
M.decode_fixed64 = decode_fixed64

-- ---------------------------------------------------------------------------
-- Float / Double (IEEE 754 little-endian)
-- ---------------------------------------------------------------------------

ffi.cdef[[
    typedef union { float    f; uint32_t u; uint8_t b[4]; } pb_f32_u_t;
    typedef union { double   d; uint64_t u; uint8_t b[8]; } pb_f64_u_t;
]]

local F32 = ffi.new('pb_f32_u_t')
local F64 = ffi.new('pb_f64_u_t')

local function encode_float(n)
    F32.f = n
    return ffi.string(F32.b, 4)
end
M.encode_float = encode_float

local function decode_float(buf, pos)
    if pos + 3 > #buf then error("truncated float", 0) end
    ffi.copy(F32.b, buf:sub(pos, pos + 3), 4)
    return tonumber(F32.f), pos + 4
end
M.decode_float = decode_float

local function encode_double(n)
    F64.d = n
    return ffi.string(F64.b, 8)
end
M.encode_double = encode_double

local function decode_double(buf, pos)
    if pos + 7 > #buf then error("truncated double", 0) end
    ffi.copy(F64.b, buf:sub(pos, pos + 7), 8)
    return tonumber(F64.d), pos + 8
end
M.decode_double = decode_double

-- ---------------------------------------------------------------------------
-- Length-delimited (LEN)
-- ---------------------------------------------------------------------------

local function encode_len(s)
    return encode_varint(#s) .. s
end
M.encode_len = encode_len

-- decode_len(buf, pos) -> string, new_pos
local function decode_len(buf, pos)
    local b = buf:byte(pos)
    if b == nil then error("truncated varint at offset " .. pos, 0) end
    local len, npos
    if b < 0x80 then
        len = b; npos = pos + 1
    else
        local v
        v, npos = decode_varint(buf, pos)
        len = tonumber(v)
    end
    if npos + len - 1 > #buf then error("truncated LEN payload", 0) end
    return buf:sub(npos, npos + len - 1), npos + len
end
M.decode_len = decode_len

-- ---------------------------------------------------------------------------
-- Typed scalar encoders/decoders (one per proto3 scalar type).
--
-- Both code paths (descriptor-driven runtime and inline codegen) call into
-- these. They take the application-side Lua value and produce/consume only
-- the value's wire bytes — the field tag is the caller's responsibility.
-- ---------------------------------------------------------------------------

-- Encoders --------------------------------------------------------------------
M.encode_int32  = encode_varint
M.encode_int64  = encode_varint
M.encode_uint32 = encode_varint
M.encode_uint64 = encode_varint

local function encode_sint32(v) return encode_varint(zigzag_encode32(v)) end
local function encode_sint64(v) return encode_varint(zigzag_encode64(v)) end
local function encode_bool(v)   return encode_varint(v and 1 or 0)       end
M.encode_sint32 = encode_sint32
M.encode_sint64 = encode_sint64
M.encode_bool   = encode_bool

M.encode_sfixed32 = encode_fixed32  -- bits are identical, only interpretation differs
M.encode_sfixed64 = encode_fixed64
M.encode_string   = encode_len
M.encode_bytes    = encode_len
-- (encode_fixed32, encode_fixed64, encode_float, encode_double already on M)

-- Decoders --------------------------------------------------------------------
-- The 1-byte fast path is inlined at every varint-based scalar decoder
-- (see comment above decode_tag). Each decoder reads the first byte,
-- handles the common 0..127 case in straight-line code, and falls
-- through to decode_varint only for multi-byte values.
local function decode_int32(buf, pos)
    local b = buf:byte(pos)
    if b == nil then error("truncated varint at offset " .. pos, 0) end
    if b < 0x80 then return b, pos + 1 end  -- 0..127 fits int32 directly
    local u, np = decode_varint(buf, pos)
    return varint_to_int32(u), np
end
local function decode_int64(buf, pos)
    local b = buf:byte(pos)
    if b == nil then error("truncated varint at offset " .. pos, 0) end
    if b < 0x80 then return INT64(b), pos + 1 end
    local u, np = decode_varint(buf, pos)
    return INT64(u), np
end
local function decode_uint32(buf, pos)
    local b = buf:byte(pos)
    if b == nil then error("truncated varint at offset " .. pos, 0) end
    if b < 0x80 then return b, pos + 1 end
    local u, np = decode_varint(buf, pos)
    return varint_to_uint32(u), np
end
local function decode_uint64(buf, pos)
    local b = buf:byte(pos)
    if b == nil then error("truncated varint at offset " .. pos, 0) end
    if b < 0x80 then return UINT64(b), pos + 1 end
    local u, np = decode_varint(buf, pos)
    return UINT64(u), np
end
local function decode_sint32(buf, pos)
    local b = buf:byte(pos)
    if b == nil then error("truncated varint at offset " .. pos, 0) end
    if b < 0x80 then return zigzag_decode32(b), pos + 1 end
    local u, np = decode_varint(buf, pos)
    return zigzag_decode32(varint_to_uint32(u)), np
end
local function decode_sint64(buf, pos)
    local b = buf:byte(pos)
    if b == nil then error("truncated varint at offset " .. pos, 0) end
    if b < 0x80 then return zigzag_decode64(b), pos + 1 end
    local u, np = decode_varint(buf, pos)
    return zigzag_decode64(u), np
end
local function decode_bool(buf, pos)
    local b = buf:byte(pos)
    if b == nil then error("truncated varint at offset " .. pos, 0) end
    if b < 0x80 then return b ~= 0, pos + 1 end
    local u, np = decode_varint(buf, pos)
    return u ~= UINT64_ZERO, np
end
local function decode_sfixed32(buf, pos)
    local n, np = decode_fixed32(buf, pos)
    if n > 0x7fffffff then n = n - 0x100000000 end
    return n, np
end
local function decode_sfixed64(buf, pos)
    local u, np = decode_fixed64(buf, pos)
    return INT64(u), np
end

M.decode_int32    = decode_int32
M.decode_int64    = decode_int64
M.decode_uint32   = decode_uint32
M.decode_uint64   = decode_uint64
M.decode_sint32   = decode_sint32
M.decode_sint64   = decode_sint64
M.decode_bool     = decode_bool
M.decode_sfixed32 = decode_sfixed32
M.decode_sfixed64 = decode_sfixed64
-- RFC 3629 UTF-8 validator. Rejects: out-of-range continuation bytes,
-- truncated multi-byte sequences, overlong encodings, UTF-16 surrogate
-- code points (U+D800..U+DFFF), and code points above U+10FFFF.
local function is_valid_utf8(s)
    local i, n = 1, #s
    while i <= n do
        local b = s:byte(i)
        if b < 0x80 then
            i = i + 1
        elseif b < 0xC2 then
            return false  -- stray continuation or overlong 2-byte
        elseif b < 0xE0 then
            if i + 1 > n then return false end
            local b2 = s:byte(i + 1)
            if b2 < 0x80 or b2 > 0xBF then return false end
            i = i + 2
        elseif b < 0xF0 then
            if i + 2 > n then return false end
            local b2 = s:byte(i + 1)
            local b3 = s:byte(i + 2)
            if b == 0xE0 and b2 < 0xA0 then return false end  -- overlong
            if b == 0xED and b2 > 0x9F then return false end  -- surrogate
            if b2 < 0x80 or b2 > 0xBF or b3 < 0x80 or b3 > 0xBF then
                return false
            end
            i = i + 3
        elseif b < 0xF5 then
            if i + 3 > n then return false end
            local b2 = s:byte(i + 1)
            local b3 = s:byte(i + 2)
            local b4 = s:byte(i + 3)
            if b == 0xF0 and b2 < 0x90 then return false end  -- overlong
            if b == 0xF4 and b2 > 0x8F then return false end  -- > U+10FFFF
            if b2 < 0x80 or b2 > 0xBF
               or b3 < 0x80 or b3 > 0xBF
               or b4 < 0x80 or b4 > 0xBF then
                return false
            end
            i = i + 4
        else
            return false  -- 5-byte+ sequences or 0xF5..0xFF
        end
    end
    return true
end
M.is_valid_utf8 = is_valid_utf8

local function decode_string(buf, pos)
    local s, np = decode_len(buf, pos)
    if not is_valid_utf8(s) then
        error("invalid UTF-8 in string field at offset " .. pos, 0)
    end
    return s, np
end
M.decode_string = decode_string
M.decode_bytes  = decode_len
-- (decode_fixed32, decode_fixed64, decode_float, decode_double already on M
-- and have the right semantics for their proto types: fixed32 -> uint32 Lua
-- number 0..2^32-1, fixed64 -> uint64 cdata.)

-- ---------------------------------------------------------------------------
-- TYPE_INFO — single-source-of-truth metadata for the codec layer and codegen.
-- Each entry carries the wire type, packed-list eligibility, and the typed
-- encode/decode functions defined above. Adapted from tarantool-etcd's
-- types.lua pattern.
-- ---------------------------------------------------------------------------

M.TYPE_INFO = {
    int32    = {wire = M.WIRE_VARINT, packable = true,  encode = M.encode_int32,    decode = M.decode_int32   },
    int64    = {wire = M.WIRE_VARINT, packable = true,  encode = M.encode_int64,    decode = M.decode_int64   },
    uint32   = {wire = M.WIRE_VARINT, packable = true,  encode = M.encode_uint32,   decode = M.decode_uint32  },
    uint64   = {wire = M.WIRE_VARINT, packable = true,  encode = M.encode_uint64,   decode = M.decode_uint64  },
    sint32   = {wire = M.WIRE_VARINT, packable = true,  encode = M.encode_sint32,   decode = M.decode_sint32  },
    sint64   = {wire = M.WIRE_VARINT, packable = true,  encode = M.encode_sint64,   decode = M.decode_sint64  },
    bool     = {wire = M.WIRE_VARINT, packable = true,  encode = M.encode_bool,     decode = M.decode_bool    },
    fixed32  = {wire = M.WIRE_I32,    packable = true,  encode = M.encode_fixed32,  decode = M.decode_fixed32 },
    sfixed32 = {wire = M.WIRE_I32,    packable = true,  encode = M.encode_sfixed32, decode = M.decode_sfixed32},
    float    = {wire = M.WIRE_I32,    packable = true,  encode = M.encode_float,    decode = M.decode_float   },
    fixed64  = {wire = M.WIRE_I64,    packable = true,  encode = M.encode_fixed64,  decode = M.decode_fixed64 },
    sfixed64 = {wire = M.WIRE_I64,    packable = true,  encode = M.encode_sfixed64, decode = M.decode_sfixed64},
    double   = {wire = M.WIRE_I64,    packable = true,  encode = M.encode_double,   decode = M.decode_double  },
    string   = {wire = M.WIRE_LEN,    packable = false, encode = M.encode_string,   decode = M.decode_string  },
    bytes    = {wire = M.WIRE_LEN,    packable = false, encode = M.encode_bytes,    decode = M.decode_bytes   },
}

-- ---------------------------------------------------------------------------
-- Skip an unknown field (used by decoder when an unrecognized id appears).
-- skip_field(buf, pos, wire_type) -> new_pos
-- ---------------------------------------------------------------------------
local function skip_field(buf, pos, wire_type)
    if wire_type == M.WIRE_VARINT then
        local b = buf:byte(pos)
        if b == nil then error("truncated varint at offset " .. pos, 0) end
        if b < 0x80 then return pos + 1 end
        local _, npos = decode_varint(buf, pos)
        return npos
    elseif wire_type == M.WIRE_I64 then
        local np = pos + 8
        if np > #buf + 1 then error("truncated I64 at offset " .. pos, 0) end
        return np
    elseif wire_type == M.WIRE_LEN then
        local b = buf:byte(pos)
        if b == nil then error("truncated varint at offset " .. pos, 0) end
        if b < 0x80 then
            local np = pos + 1 + b
            if np > #buf + 1 then error("truncated LEN at offset " .. pos, 0) end
            return np
        end
        local len, npos = decode_varint(buf, pos)
        local np = npos + tonumber(len)
        if np > #buf + 1 then error("truncated LEN at offset " .. pos, 0) end
        return np
    elseif wire_type == M.WIRE_I32 then
        local np = pos + 4
        if np > #buf + 1 then error("truncated I32 at offset " .. pos, 0) end
        return np
    end
    error("unknown wire type " .. tostring(wire_type), 0)
end
M.skip_field = skip_field

return M