-- proto3 JSON encoding (canonical mapping).
--
-- Spec: https://protobuf.dev/programming-guides/proto3/#json
--
-- Highlights of how we map types both ways:
--   * 32-bit ints / float / double / bool -> JSON number/bool
--   * 64-bit ints (int64/uint64/sint64/fixed64/sfixed64) -> JSON string
--     (JSON doubles lose precision past 2^53; the spec mandates strings)
--   * bytes -> base64 string
--   * enum -> string name when known, else number
--   * message -> JSON object (camelCase keys)
--   * map<K,V> -> JSON object (keys stringified per spec)
--   * Timestamp -> ISO 8601 "YYYY-MM-DDTHH:MM:SS[.nnnnnnnnn]Z"
--   * Duration -> "<seconds>s" (decimal seconds with up to 9 fractional digits)
--   * Empty -> {}
--   * Wrapper messages -> unwrapped scalar
--
-- Field names: emitted camelCase per spec; decoder accepts both camelCase
-- and the original snake_case so users aren't punished for either convention.
local ffi      = require('ffi')
local json     = require('json')
local digest   = require('digest')
local datetime = require('datetime')
local wire     = require('pb.wire')
local pbwkt    = require('pb.wkt')

local M = {}

local TYPE_INFO = wire.TYPE_INFO
local INT64_FAMILY = {int64=true, uint64=true, sint64=true, fixed64=true, sfixed64=true}

-- ---------------------------------------------------------------------------
-- Helpers
-- ---------------------------------------------------------------------------

-- snake_case -> camelCase (proto field naming convention for JSON).
local function to_camel(name)
    return (name:gsub('_(%w)', function(c) return c:upper() end))
end

-- Lossless stringification of an integer/cdata for JSON output.
local function int_to_string(v)
    if type(v) == 'cdata' then
        return tostring(v):gsub('U?LL$', '')
    end
    return tostring(v)
end

-- Convert a JSON string/number back to int64 or uint64 cdata.
local function string_to_int64(v, is_unsigned)
    if type(v) == 'number' then v = string.format('%.0f', v) end
    if type(v) ~= 'string' then error('expected JSON string or number for int64', 0) end
    local cdata = tonumber64(v)
    if cdata == nil then error('invalid int64 string: ' .. v, 0) end
    if is_unsigned then return ffi.cast('uint64_t', cdata) end
    return ffi.cast('int64_t', cdata)
end

-- ---------------------------------------------------------------------------
-- Encode (proto-Lua table -> Lua table suitable for json.encode)
-- ---------------------------------------------------------------------------

local to_json_value  -- forward
local encode_message -- forward

-- Encode a single scalar (not repeated/map). Returns a JSON-friendly value.
local function encode_scalar(proto_type, v)
    if INT64_FAMILY[proto_type] then return int_to_string(v) end
    if proto_type == 'uint32' then
        -- Lua double can hold 0..2^32 - 1; emit as number.
        return v
    end
    if proto_type == 'bytes' then return digest.base64_encode(v) end
    if proto_type == 'float' or proto_type == 'double' then
        -- JSON has no NaN/Infinity literals; spec mandates string sentinels.
        if v ~= v then return 'NaN' end
        if v == math.huge then return 'Infinity' end
        if v == -math.huge then return '-Infinity' end
        return v
    end
    return v  -- string, bool, int32, sint32, fixed32, sfixed32
end

local function encode_enum(enum_desc, v)
    if type(v) == 'string' then return v end
    local name = enum_desc.by_value[v]
    return name or v  -- unknown numeric value: emit as number
end

-- Convert a single proto value to a JSON-encodable value, based on field kind.
local function encode_field_value(field, v)
    local kind = field.kind
    if kind == 'scalar' then
        return encode_scalar(field.proto_type, v)
    elseif kind == 'enum' then
        return encode_enum(field.enum, v)
    elseif kind == 'message' then
        return encode_message(field.message, v)
    end
    error('encode_field_value: unknown kind ' .. tostring(kind), 0)
end

-- Map key encoding: per spec, all keys are strings in JSON output.
local function encode_map_key(key_field, k)
    local pt = key_field.proto_type
    if pt == 'bool' then return k and 'true' or 'false' end
    if INT64_FAMILY[pt] then return int_to_string(k) end
    if pt == 'string' then return k end
    return tostring(k)  -- int32/uint32/sint32/etc.
end

-- Struct/Value/ListValue JSON mappings.
--
-- Per the proto3 JSON spec:
--   * Struct    ↔ JSON object  (the `fields` map is hoisted away)
--   * ListValue ↔ JSON array   (the `values` array is hoisted away)
--   * Value     ↔ any JSON value (null/number/string/bool/object/array)
--
-- Lua representations mirror what runtime/pb/wkt.lua produces:
--   * box.NULL → JSON null
--   * boolean/number/string → JSON true|false/number/string
--   * table tagged with pb.wkt.list → JSON array; otherwise → JSON object

local PB_NULL = pbwkt.NULL

local value_to_json, value_to_json_struct, value_to_json_list  -- forwards

value_to_json = function(v)
    if v == nil or v == PB_NULL then return box.NULL end
    local ty = type(v)
    if ty == 'boolean' or ty == 'number' or ty == 'string' then return v end
    if ty == 'cdata' then return tonumber(v) end
    if ty == 'table' then
        local mt = getmetatable(v)
        if mt and mt.__pb_kind == 'list'   then return value_to_json_list(v)   end
        if mt and mt.__pb_kind == 'struct' then return value_to_json_struct(v) end
        if v[1] ~= nil then return value_to_json_list(v) end
        return value_to_json_struct(v)
    end
    error('Value JSON: unsupported Lua type ' .. ty, 0)
end

value_to_json_struct = function(t)
    if t == nil then return setmetatable({}, {__serialize='map'}) end
    local out, empty = {}, true
    for k, v in pairs(t) do
        out[tostring(k)] = value_to_json(v)
        empty = false
    end
    if empty then return setmetatable(out, {__serialize='map'}) end
    return out
end

value_to_json_list = function(t)
    if t == nil then return setmetatable({}, {__serialize='seq'}) end
    local out = {}
    for i = 1, #t do out[i] = value_to_json(t[i]) end
    return setmetatable(out, {__serialize='seq'})
end

-- FieldMask JSON: paths joined by `,`. snake_case → lowerCamelCase per spec.
local function fieldmask_to_json(v)
    if v == nil or #v == 0 then return '' end
    local parts = {}
    for i = 1, #v do parts[i] = to_camel(v[i]) end
    return table.concat(parts, ',')
end

local function fieldmask_from_json(s)
    if type(s) ~= 'string' or s == '' then return {} end
    local out = {}
    for part in (s .. ','):gmatch('([^,]+),') do
        -- lowerCamelCase -> snake_case
        out[#out + 1] = (part:gsub('(%u)', function(c) return '_' .. c:lower() end))
    end
    return out
end

-- Any JSON: a flat object {"@type": "<url>", ...fields...}. Decoded by
-- consulting pb.wkt registry. Pack/unpack the payload through the registered
-- descriptor; unregistered types fall back to the opaque {type_url,value} form.
local function any_to_json(v)
    if v == nil then return setmetatable({}, {__serialize='map'}) end
    if type(v) ~= 'table' then
        error('Any JSON: expected table, got ' .. type(v), 0)
    end
    local type_url = v.type_url or ''
    local bytes    = v.value or ''
    local desc     = pbwkt.lookup(type_url)
    if desc == nil or bytes == '' then
        -- Opaque fallback: emit the protobuf representation as-is so a round
        -- trip is still possible without a registered descriptor.
        local obj = {['@type'] = type_url}
        if bytes ~= '' then obj.value = digest.base64_encode(bytes) end
        return obj
    end
    local inner = desc.decode and desc.decode(bytes)
                or require('pb.codec').decode(desc, bytes)
    local payload = encode_message(desc, inner)
    -- For Value/Struct/ListValue/wrappers the JSON form is not an object; the
    -- spec says to nest under "value" then.
    if type(payload) ~= 'table' or getmetatable(payload) and
       getmetatable(payload).__serialize == 'seq' then
        return {['@type'] = type_url, value = payload}
    end
    payload['@type'] = type_url
    return payload
end

-- WKT special-case encoders. Return either a JSON-encodable Lua value or
-- nil to indicate "no override; fall back to generic message walk".
local function encode_wkt(desc, v)
    local name = desc.name
    if name == 'google.protobuf.Empty' then
        return setmetatable({}, {__serialize='map'})  -- emit as {}
    end
    if name == 'google.protobuf.Timestamp' then
        local dt = v
        if type(v) == 'table' then
            dt = datetime.new({timestamp = tonumber(v.seconds or 0),
                               nsec      = v.nanos or 0})
        elseif type(v) == 'cdata' and datetime.is_datetime(v) then
            dt = v
        elseif type(v) == 'number' then
            dt = datetime.new({timestamp = math.floor(v),
                               nsec      = math.floor((v - math.floor(v)) * 1e9 + 0.5)})
        else
            error('Timestamp JSON: expected datetime/table/number', 0)
        end
        return tostring(dt):gsub('Z$', 'Z')  -- Tarantool already emits ISO 8601
    end
    if name == 'google.protobuf.Duration' then
        local seconds, nanos
        if type(v) == 'table' then
            seconds = tonumber(v.seconds or 0)
            nanos = v.nanos or 0
        elseif type(v) == 'number' then
            seconds = math.floor(v)
            nanos = math.floor((v - seconds) * 1e9 + 0.5)
        else
            error('Duration JSON: expected table or number', 0)
        end
        local total_secs = seconds
        if nanos == 0 then return string.format('%ds', total_secs) end
        local fraction = string.format('.%09d', nanos):gsub('0+$', '')
        if fraction == '.' then fraction = '' end
        return string.format('%d%ss', total_secs, fraction)
    end
    -- Wrappers: encode is just the unwrapped value.
    local wrap = name:match('^google%.protobuf%.(%w+)Value$')
    if wrap then
        local wrapper_proto = {
            Int32 = 'int32', UInt32 = 'uint32', Int64 = 'int64', UInt64 = 'uint64',
            Float = 'float', Double = 'double', Bool = 'bool',
            String = 'string', Bytes = 'bytes',
        }
        local pt = wrapper_proto[wrap]
        if pt then return encode_scalar(pt, v) end
    end
    if name == 'google.protobuf.Struct' then
        return value_to_json_struct(v)
    end
    if name == 'google.protobuf.ListValue' then
        return value_to_json_list(v)
    end
    if name == 'google.protobuf.Value' then
        return value_to_json(v)
    end
    if name == 'google.protobuf.FieldMask' then
        return fieldmask_to_json(v)
    end
    if name == 'google.protobuf.Any' then
        return any_to_json(v)
    end
    return nil
end

encode_message = function(desc, t)
    if t == nil then return nil end
    local override = encode_wkt(desc, t)
    if override ~= nil then return override end

    local out = {}
    for _, f in ipairs(desc.fields) do
        local v = t[f.name]
        if v ~= nil then
            local key = to_camel(f.name)
            if f.kind == 'map' then
                if next(v) ~= nil then
                    local obj = {}
                    for k, mv in pairs(v) do
                        obj[encode_map_key(f.key, k)] = encode_field_value(f.value, mv)
                    end
                    out[key] = obj
                end
            elseif f.repeated then
                if #v > 0 then
                    local arr = {}
                    for i = 1, #v do arr[i] = encode_field_value(f, v[i]) end
                    out[key] = arr
                end
            else
                -- proto3 default elision (unless presence is meaningful).
                local emit = true
                if not (f.optional or f.oneof) then
                    if f.kind == 'scalar' then
                        local pt = f.proto_type
                        if pt == 'string' or pt == 'bytes' then
                            if v == '' then emit = false end
                        elseif pt == 'bool' then
                            if v == false then emit = false end
                        else
                            -- numeric default
                            if v == 0 or (type(v) == 'cdata' and v == ffi.cast('int64_t', 0)) then
                                emit = false
                            end
                        end
                    elseif f.kind == 'enum' then
                        if v == 0 or v == f.enum.by_value[0] then emit = false end
                    end
                end
                if emit then out[key] = encode_field_value(f, v) end
            end
        end
    end
    return out
end

to_json_value = encode_message

function M.encode(desc, t)
    return json.encode(encode_message(desc, t))
end

-- ---------------------------------------------------------------------------
-- Decode (JSON -> proto-Lua table)
-- ---------------------------------------------------------------------------

local decode_message  -- forward

local function decode_scalar(proto_type, v)
    if INT64_FAMILY[proto_type] then
        return string_to_int64(v, proto_type:sub(1, 1) == 'u' or proto_type == 'fixed64')
    end
    if proto_type == 'bytes' then return digest.base64_decode(v) end
    if proto_type == 'float' or proto_type == 'double' then
        if v == 'NaN' then return 0/0 end
        if v == 'Infinity' then return math.huge end
        if v == '-Infinity' then return -math.huge end
        if type(v) == 'string' then return tonumber(v) end
        return v
    end
    if proto_type == 'bool' then
        if type(v) == 'string' then return v == 'true' end
        return v and true or false
    end
    -- 32-bit ints + string: accept JSON string or number defensively.
    if type(v) == 'string' and proto_type ~= 'string' then return tonumber(v) end
    return v
end

local function decode_enum(enum_desc, v)
    if type(v) == 'string' then
        local n = enum_desc.by_name[v]
        if n ~= nil then return n end
    end
    -- proto3 JSON: unrecognized integer enum values pass through; unrecognized
    -- string names yield nil so the caller can drop the element (repeated/map)
    -- or fall back to the field default (singular).
    return tonumber(v)
end

local function decode_field_value(field, v)
    local kind = field.kind
    if kind == 'scalar' then return decode_scalar(field.proto_type, v)
    elseif kind == 'enum' then return decode_enum(field.enum, v)
    elseif kind == 'message' then return decode_message(field.message, v) end
    error('decode_field_value: unknown kind ' .. tostring(kind), 0)
end

local function decode_map_key(key_field, k)
    local pt = key_field.proto_type
    if pt == 'string' then return k end
    if pt == 'bool' then return k == 'true' end
    if INT64_FAMILY[pt] then return string_to_int64(k, pt:sub(1, 1) == 'u' or pt == 'fixed64') end
    return tonumber(k)
end

local json_to_value, json_to_struct, json_to_list, json_to_any  -- forwards

json_to_any = function(v)
    if v == nil then return {type_url = '', value = ''} end
    if type(v) ~= 'table' then
        error('Any JSON: expected object, got ' .. type(v), 0)
    end
    local type_url = v['@type'] or ''
    local desc = pbwkt.lookup(type_url)
    if desc == nil then
        -- Opaque fallback (no registered descriptor): expect a base64 `value`
        -- field, mirroring our encode-side fallback.
        local raw = v.value
        return {type_url = type_url,
                value = raw and digest.base64_decode(raw) or ''}
    end
    -- Reconstruct the inner message from the flat JSON, skipping @type.
    local payload = {}
    for k, mv in pairs(v) do
        if k ~= '@type' then payload[k] = mv end
    end
    -- For Value/Struct/ListValue/wrappers the spec nests under `value`.
    if payload.value ~= nil and next(payload, next(payload)) == nil then
        payload = payload.value
    end
    local inner = decode_message(desc, payload)
    local bytes = desc.encode and desc.encode(inner)
                or require('pb.codec').encode(desc, inner)
    return {type_url = type_url, value = bytes}
end


json_to_value = function(v)
    if v == nil or v == box.NULL then return PB_NULL end
    local ty = type(v)
    if ty == 'boolean' or ty == 'number' or ty == 'string' then return v end
    if ty == 'cdata' then return tonumber(v) end
    if ty == 'table' then
        if v[1] ~= nil or next(v) == nil and getmetatable(v)
                       and getmetatable(v).__serialize == 'seq' then
            return json_to_list(v)
        end
        -- Heuristic: integer-keyed → list, else → struct.
        if v[1] ~= nil then return json_to_list(v) end
        return json_to_struct(v)
    end
    error('Value JSON: unsupported type ' .. ty, 0)
end

json_to_struct = function(v)
    local out = pbwkt.struct({})
    if v == nil then return out end
    for k, mv in pairs(v) do out[k] = json_to_value(mv) end
    return out
end

json_to_list = function(v)
    local out = pbwkt.list({})
    if v == nil then return out end
    for i = 1, #v do out[i] = json_to_value(v[i]) end
    return out
end

local function decode_wkt(desc, v)
    local name = desc.name
    if name == 'google.protobuf.Empty' then return {} end
    if name == 'google.protobuf.Timestamp' then
        return datetime.parse(v)
    end
    if name == 'google.protobuf.Duration' then
        local body = v:gsub('s$', '')
        local sign = 1
        if body:sub(1, 1) == '-' then sign = -1; body = body:sub(2) end
        local sec_str, frac = body:match('^(%d+)%.?(%d*)$')
        if sec_str == nil then error('invalid Duration JSON: ' .. v, 0) end
        local nanos = 0
        if frac and frac ~= '' then
            nanos = tonumber((frac .. '000000000'):sub(1, 9))
        end
        return {seconds = sign * tonumber(sec_str), nanos = sign * nanos}
    end
    local wrap = name:match('^google%.protobuf%.(%w+)Value$')
    if wrap then
        local wrapper_proto = {
            Int32 = 'int32', UInt32 = 'uint32', Int64 = 'int64', UInt64 = 'uint64',
            Float = 'float', Double = 'double', Bool = 'bool',
            String = 'string', Bytes = 'bytes',
        }
        local pt = wrapper_proto[wrap]
        if pt then return decode_scalar(pt, v) end
    end
    if name == 'google.protobuf.Value' then
        return json_to_value(v)
    end
    if name == 'google.protobuf.Struct' then
        return json_to_struct(v)
    end
    if name == 'google.protobuf.ListValue' then
        return json_to_list(v)
    end
    if name == 'google.protobuf.FieldMask' then
        return fieldmask_from_json(v)
    end
    if name == 'google.protobuf.Any' then
        return json_to_any(v)
    end
    return nil
end

decode_message = function(desc, v)
    if v == nil then return nil end
    local override = decode_wkt(desc, v)
    if override ~= nil then return override end
    if type(v) ~= 'table' then
        error('expected JSON object for ' .. desc.name .. ', got ' .. type(v), 0)
    end

    -- Build a name -> field map covering both camelCase and snake_case.
    local field_by_json_name = desc._json_field_by_name
    if field_by_json_name == nil then
        field_by_json_name = {}
        for _, f in ipairs(desc.fields) do
            field_by_json_name[f.name] = f
            field_by_json_name[to_camel(f.name)] = f
        end
        desc._json_field_by_name = field_by_json_name
    end

    local out = {}
    for k, jv in pairs(v) do
        local f = field_by_json_name[k]
        if f ~= nil then
            if f.kind == 'map' then
                local m = {}
                for mk, mv in pairs(jv) do
                    local dv = decode_field_value(f.value, mv)
                    if dv ~= nil then
                        m[decode_map_key(f.key, mk)] = dv
                    end
                end
                out[f.name] = m
            elseif f.repeated then
                local arr = {}
                local n = 0
                for i = 1, #jv do
                    local dv = decode_field_value(f, jv[i])
                    if dv ~= nil then n = n + 1; arr[n] = dv end
                end
                out[f.name] = arr
            else
                local dv = decode_field_value(f, jv)
                if dv ~= nil then out[f.name] = dv end
            end
        end
        -- Unknown JSON keys are silently ignored (per spec).
    end
    return out
end

function M.decode(desc, s)
    return decode_message(desc, json.decode(s))
end

M.to_json_value   = to_json_value
M.from_json_value = decode_message

return M