-- Lazy, zero-copy decode view over protobuf wire bytes.
--
-- Built once via pb.decode_lazy(desc, bytes). Indexes the wire bytes at
-- construct time (single pass), then materializes individual fields on
-- access. Nested messages, repeated fields, and maps are themselves
-- returned as lazy sub-views; values are decoded on `:get` / `:at`.
--
-- The underlying Lua string is kept GC-anchored on the view table so
-- substring-based passthrough remains valid for the view's lifetime.
--
-- Phase 1: read-only API (:get / :has / :which / :iter / :names plus
-- ArrayView :len/:at/:iter and MapView :get/:has/:keys/:iter). Mutation
-- and passthrough re-encode arrive in phase 2.
--
-- WKT and any descriptor carrying `desc.decode` are eagerly decoded and
-- wrapped in an EagerView with the same getter surface, so callers don't
-- have to special-case Timestamp/Duration/Struct/etc.

local wire  = require('pb.wire')
local codec = require('pb.codec')

local M = {}

-- Forward decls: views can reference each other.
local build_msg_view
local build_array_view
local build_map_view

-- ---------------------------------------------------------------------------
-- Index build
--
-- Single pass over `bytes`. Each tag+value occupies one `segment`:
--   {id = <field id>,
--    tag_start = <1-based offset of tag's first byte>,
--    val_start = <1-based offset of value's first byte>,
--    next_start = <1-based offset just past this segment>,
--    wt = <wire type>}
--
-- `segments` is the wire-order list (used by iter and by passthrough re-encode
-- in phase 2). `by_id` maps id -> list-of-segments, even for singular fields
-- (multiple wire entries for the same singular id are legal per spec: scalar
-- last-wins, message-merge — both consumers can walk the list).
-- ---------------------------------------------------------------------------
local function index_bytes(desc, bytes)
    local pos, lim = 1, #bytes
    local segments = {}
    local by_id = {}
    local fbi = desc.field_by_id

    while pos <= lim do
        local tag_start = pos
        local id, wt, npos = wire.decode_tag(bytes, pos)
        local val_start = npos
        local next_start = wire.skip_field(bytes, npos, wt)
        local seg = {
            id = id,
            tag_start = tag_start,
            val_start = val_start,
            next_start = next_start,
            wt = wt,
        }
        segments[#segments + 1] = seg
        if fbi[id] ~= nil then
            local list = by_id[id]
            if list == nil then
                by_id[id] = {seg}
            else
                list[#list + 1] = seg
            end
        end
        pos = next_start
    end
    return segments, by_id
end

-- ---------------------------------------------------------------------------
-- Per-field materialization
-- ---------------------------------------------------------------------------

-- Decode a single value from one segment, given the singular field shape.
-- For message fields this returns another lazy view; for scalar/enum it
-- returns the materialized Lua/cdata value.
local function read_singular(field, bytes, seg)
    local kind = field.kind
    if kind == 'scalar' then
        local h = wire.TYPE_INFO[field.proto_type]
        local v = h.decode(bytes, seg.val_start)
        return v
    elseif kind == 'enum' then
        local u = wire.decode_varint(bytes, seg.val_start)
        return tonumber(u)
    elseif kind == 'message' then
        local payload = wire.decode_len(bytes, seg.val_start)
        if field.message.decode ~= nil then
            -- WKT (or any custom-decode override). Eager.
            return field.message.decode(payload)
        end
        return build_msg_view(field.message, payload)
    end
    error("read_singular: unknown kind " .. tostring(kind), 0)
end

-- For singular fields, semantics for multiple wire entries with the same id:
--   - scalar/enum: last wins.
--   - message: merged. We delegate to the eager codec by concatenating the
--     per-entry payloads and feeding them to pb.codec.decode, which honors
--     proto3 merge rules. (Multi-entry singular messages are rare; this is
--     the off-fast-path correctness branch.)
local function read_singular_list(field, bytes, list)
    if field.kind ~= 'message' or #list == 1 then
        return read_singular(field, bytes, list[#list])
    end
    -- Multi-segment message: concatenate the inner payloads and eager-decode.
    -- This loses the lazy sub-view but is the right thing semantically.
    local parts = {}
    for i = 1, #list do
        local seg = list[i]
        local payload = wire.decode_len(bytes, seg.val_start)
        parts[i] = payload
    end
    local merged = table.concat(parts)
    if field.message.decode ~= nil then
        return field.message.decode(merged)
    end
    return build_msg_view(field.message, merged)
end

-- ---------------------------------------------------------------------------
-- ArrayView: lazy view over a repeated field.
--
-- For unpacked repeated, each wire entry is one element (one segment).
-- For packed repeated (scalars/enums), one segment contains a length-delimited
-- payload with all elements; we walk the payload on demand. We never
-- materialize the full element array unless the user iterates it all.
-- ---------------------------------------------------------------------------

local ArrayView = {}
ArrayView.__index = ArrayView

-- Pre-expand packed payloads into a uniform per-element segment list at
-- construct time. Cheaper to do once than to re-scan on every :at(i).
local function expand_packed(field, bytes, packed_seg)
    local h = field.kind == 'scalar' and wire.TYPE_INFO[field.proto_type] or nil
    -- Decode the length prefix to know where elements live.
    local payload_start = packed_seg.val_start
    local b = bytes:byte(payload_start)
    local payload_len, hdr_end
    if b < 0x80 then
        payload_len = b
        hdr_end = payload_start + 1
    else
        local v, npos = wire.decode_varint(bytes, payload_start)
        payload_len = tonumber(v)
        hdr_end = npos
    end
    local lim = hdr_end + payload_len
    local elems = {}
    local p = hdr_end
    while p < lim do
        local v_start = p
        local v_next
        if field.kind == 'scalar' then
            -- Skip according to the scalar's actual wire type.
            v_next = wire.skip_field(bytes, p, h.wire)
        else
            -- Packed enums: varint per element.
            v_next = wire.skip_field(bytes, p, wire.WIRE_VARINT)
        end
        elems[#elems + 1] = {val_start = v_start, val_next = v_next}
        p = v_next
    end
    return elems
end

local function build_array_view_impl(field, bytes, segs)
    -- For packed payloads (single segment, WIRE_LEN, but scalar field uses
    -- a non-LEN wire type), elements live inside that one segment.
    -- For unpacked, each segment is one element.
    local mode  -- 'packed' | 'unpacked'
    local elements  -- list of {val_start = ...} entries
    if field.kind == 'scalar' then
        local h = wire.TYPE_INFO[field.proto_type]
        if h.wire ~= wire.WIRE_LEN and #segs == 1 and segs[1].wt == wire.WIRE_LEN then
            mode = 'packed'
            elements = expand_packed(field, bytes, segs[1])
        else
            mode = 'unpacked'
            elements = segs
        end
    elseif field.kind == 'enum' then
        if #segs == 1 and segs[1].wt == wire.WIRE_LEN then
            mode = 'packed'
            elements = expand_packed(field, bytes, segs[1])
        else
            mode = 'unpacked'
            elements = segs
        end
    else
        -- repeated message: always one segment per element, never packed.
        mode = 'unpacked'
        elements = segs
    end
    return setmetatable({
        _field = field,
        _bytes = bytes,
        _mode  = mode,
        _elements = elements,
    }, ArrayView)
end

build_array_view = build_array_view_impl

function ArrayView:len()
    return #self._elements
end

function ArrayView:at(i)
    local entry = self._elements[i]
    if entry == nil then return nil end
    local field, bytes = self._field, self._bytes
    if self._mode == 'packed' then
        if field.kind == 'enum' then
            local u = wire.decode_varint(bytes, entry.val_start)
            return tonumber(u)
        end
        local h = wire.TYPE_INFO[field.proto_type]
        return (h.decode(bytes, entry.val_start))
    end
    return read_singular(field, bytes, entry)
end

function ArrayView:iter()
    local view, i = self, 0
    local n = #self._elements
    return function()
        i = i + 1
        if i > n then return nil end
        return i, view:at(i)
    end
end

function ArrayView:tolist()
    local out = {}
    for i = 1, #self._elements do out[i] = self:at(i) end
    return out
end

-- ---------------------------------------------------------------------------
-- MapView: lazy view over a map<K,V> field.
--
-- Each wire entry is one length-delimited submessage with two fields:
--   id=1 (key), id=2 (value). The submessage may omit either when the
-- value equals its proto3 default.
--
-- For O(1) `:get(k)`, we decode keys lazily but cache the key->entry map on
-- first key-lookup or first :iter. Until then, only the per-entry payload
-- offsets are known.
-- ---------------------------------------------------------------------------

local MapView = {}
MapView.__index = MapView

local function decode_map_entry(field, bytes, seg)
    local key_field, val_field = field.key, field.value
    -- Read the outer LEN to find the entry payload.
    local b = bytes:byte(seg.val_start)
    local payload_len, hdr_end
    if b < 0x80 then
        payload_len = b; hdr_end = seg.val_start + 1
    else
        local v, npos = wire.decode_varint(bytes, seg.val_start)
        payload_len = tonumber(v); hdr_end = npos
    end
    local lim = hdr_end + payload_len
    local key, val
    local p = hdr_end
    while p < lim do
        local eid, ewt, np = wire.decode_tag(bytes, p)
        p = np
        if eid == 1 then
            if key_field.kind == 'scalar' then
                key, p = wire.TYPE_INFO[key_field.proto_type].decode(bytes, p)
            elseif key_field.kind == 'enum' then
                local u; u, p = wire.decode_varint(bytes, p); key = tonumber(u)
            end
        elseif eid == 2 then
            if val_field.kind == 'scalar' then
                val, p = wire.TYPE_INFO[val_field.proto_type].decode(bytes, p)
            elseif val_field.kind == 'enum' then
                local u; u, p = wire.decode_varint(bytes, p); val = tonumber(u)
            elseif val_field.kind == 'message' then
                local payload; payload, p = wire.decode_len(bytes, p)
                if val_field.message.decode ~= nil then
                    val = val_field.message.decode(payload)
                else
                    val = build_msg_view(val_field.message, payload)
                end
            end
        else
            p = wire.skip_field(bytes, p, ewt)
        end
    end
    if key == nil then
        if key_field.kind == 'scalar' then
            local pt = key_field.proto_type
            if pt == 'string' or pt == 'bytes' then key = ''
            elseif pt == 'bool' then key = false
            else key = 0 end
        else key = 0 end
    end
    if val == nil then
        if val_field.kind == 'scalar' then
            local pt = val_field.proto_type
            if pt == 'string' or pt == 'bytes' then val = ''
            elseif pt == 'bool' then val = false
            else val = 0 end
        elseif val_field.kind == 'enum' then val = 0
        elseif val_field.kind == 'message' then val = {} end
    end
    return key, val
end

local function build_map_view_impl(field, bytes, segs)
    return setmetatable({
        _field = field,
        _bytes = bytes,
        _segs  = segs,
        -- _by_key populated lazily on first :get/:has/:iter call.
    }, MapView)
end

build_map_view = build_map_view_impl

local function map_ensure_index(self)
    if self._by_key ~= nil then return end
    local field, bytes, segs = self._field, self._bytes, self._segs
    local by_key = {}
    local keys = {}
    for i = 1, #segs do
        local k, v = decode_map_entry(field, bytes, segs[i])
        if by_key[k] == nil then keys[#keys + 1] = k end
        by_key[k] = v  -- duplicate keys: last wins (matches eager decode)
    end
    self._by_key = by_key
    self._keys = keys
end

function MapView:get(k)
    map_ensure_index(self)
    return self._by_key[k]
end

function MapView:has(k)
    map_ensure_index(self)
    return self._by_key[k] ~= nil
end

function MapView:keys()
    map_ensure_index(self)
    local out = {}
    for i = 1, #self._keys do out[i] = self._keys[i] end
    return out
end

function MapView:iter()
    map_ensure_index(self)
    local keys, by_key = self._keys, self._by_key
    local i = 0
    return function()
        i = i + 1
        local k = keys[i]
        if k == nil then return nil end
        return k, by_key[k]
    end
end

function MapView:totable()
    map_ensure_index(self)
    local out = {}
    for k, v in pairs(self._by_key) do out[k] = v end
    return out
end

-- ---------------------------------------------------------------------------
-- MessageView (top-level)
-- ---------------------------------------------------------------------------

local MessageView = {}
MessageView.__index = MessageView

local function build_msg_view_impl(desc, bytes)
    if desc.decode ~= nil then
        -- WKT / custom-decode descriptor. Eager-wrap so the API stays uniform.
        local materialized = desc.decode(bytes)
        return setmetatable({
            _desc = desc,
            _eager = materialized,
            _eager_only = true,
        }, MessageView)
    end
    local segments, by_id = index_bytes(desc, bytes)
    return setmetatable({
        _desc = desc,
        _bytes = bytes,
        _segments = segments,
        _by_id = by_id,
        _cache = {},
        -- Parallel array of cached sub-MessageViews so :is_dirty can
        -- ipairs over it instead of pairs(_cache) — `pairs` over a hash
        -- compiles to bytecode ISNEXT, NYI in Tarantool's LuaJIT 2.1.
        _sub_msg_views = {},
    }, MessageView)
end

build_msg_view = build_msg_view_impl
M.build = build_msg_view_impl

-- :get(name) -> decoded value, or nil if not on wire (and no _cache entry).
function MessageView:get(name)
    if self._eager_only then return self._eager[name] end
    local cache = self._cache
    local v = cache[name]
    if v ~= nil then return v end
    local field = self._desc.field_by_name[name]
    if field == nil then return nil end
    local segs = self._by_id[field.id]
    if segs == nil then return nil end
    if field.kind == 'map' then
        v = build_map_view(field, self._bytes, segs)
    elseif field.repeated then
        v = build_array_view(field, self._bytes, segs)
    else
        v = read_singular_list(field, self._bytes, segs)
    end
    cache[name] = v
    -- Track sub-MessageViews on a flat array so :is_dirty can walk
    -- with ipairs (see _sub_msg_views note above).
    if type(v) == 'table' and getmetatable(v) == MessageView then
        local s = self._sub_msg_views
        s[#s + 1] = v
    end
    return v
end

-- :has(name) -> was this field present on the wire?
function MessageView:has(name)
    if self._eager_only then return self._eager[name] ~= nil end
    local field = self._desc.field_by_name[name]
    if field == nil then return false end
    return self._by_id[field.id] ~= nil
end

-- :which(oneof_name) -> name of the active branch, or nil.
-- Proto3 last-wins: if multiple branches appeared on the wire, the one
-- whose final segment came last in wire order is active.
function MessageView:which(oneof_name)
    if self._eager_only then
        local oneofs = self._desc.oneofs
        if oneofs == nil then return nil end
        local members = oneofs[oneof_name]
        if members == nil then return nil end
        for i = 1, #members do
            if self._eager[members[i]] ~= nil then return members[i] end
        end
        return nil
    end
    local oneofs = self._desc.oneofs
    if oneofs == nil then return nil end
    local members = oneofs[oneof_name]
    if members == nil then return nil end
    local member_set = {}
    for i = 1, #members do member_set[members[i]] = true end
    local fbn = self._desc.field_by_name
    local active, active_pos
    local segs = self._segments
    for i = 1, #segs do
        local f = self._desc.field_by_id[segs[i].id]
        if f ~= nil and member_set[f.name] then
            if active_pos == nil or segs[i].tag_start > active_pos then
                active = f.name
                active_pos = segs[i].tag_start
            end
        end
    end
    return active
end

-- :names() -> iterator yielding present field names in wire order
-- (deduplicated; each field appears once even when it has multiple wire entries).
function MessageView:names()
    if self._eager_only then
        local fields, i = self._desc.fields, 0
        local eager = self._eager
        return function()
            while true do
                i = i + 1
                local f = fields[i]
                if f == nil then return nil end
                if eager[f.name] ~= nil then return f.name end
            end
        end
    end
    local segments = self._segments
    local fbi = self._desc.field_by_id
    local emitted = {}
    local i = 0
    return function()
        while true do
            i = i + 1
            local seg = segments[i]
            if seg == nil then return nil end
            local f = fbi[seg.id]
            if f ~= nil and emitted[f.name] == nil then
                emitted[f.name] = true
                return f.name
            end
        end
    end
end

-- :iter() -> iterator yielding (name, value) for present fields,
-- decoding each value on demand. Order = wire order, deduplicated.
function MessageView:iter()
    if self._eager_only then
        local fields, i = self._desc.fields, 0
        local eager = self._eager
        return function()
            while true do
                i = i + 1
                local f = fields[i]
                if f == nil then return nil end
                local v = eager[f.name]
                if v ~= nil then return f.name, v end
            end
        end
    end
    local segments = self._segments
    local fbi = self._desc.field_by_id
    local emitted = {}
    local view = self
    local i = 0
    return function()
        while true do
            i = i + 1
            local seg = segments[i]
            if seg == nil then return nil end
            local f = fbi[seg.id]
            if f ~= nil and emitted[f.name] == nil then
                emitted[f.name] = true
                return f.name, view:get(f.name)
            end
        end
    end
end

-- ---------------------------------------------------------------------------
-- Mutation
-- ---------------------------------------------------------------------------

-- :set(name, value) marks a field dirty. Subsequent :encode() emits the
-- new value via the codec; other fields passthrough their original bytes.
-- The provided value can be any shape the eager encoder accepts (Lua
-- table/string/number/cdata/array). Passing a MessageView as `value` is
-- supported but materializes it via :totable() on encode.
function MessageView:set(name, value)
    if self._eager_only then
        self._eager[name] = value
        self._eager_dirty = true
        return
    end
    local field = self._desc.field_by_name[name]
    if field == nil then
        error("unknown field '" .. tostring(name) .. "' on " .. self._desc.name, 0)
    end
    if self._dirty == nil then self._dirty = {} end
    self._dirty[name] = true
    self._cache[name] = value
end

-- :is_dirty() returns true if this view has had :set called, OR if any
-- cached sub-view (e.g. a nested MessageView accessed via :get) has been
-- mutated. Used by parent :encode to decide whether to splice or re-encode.
--
-- We walk _sub_msg_views (a flat array of cached sub-MessageViews) with
-- ipairs instead of pairs(_cache). Hash-keyed pairs compiles to ISNEXT
-- which is NYI in Tarantool's LuaJIT 2.1 — keeping is_dirty on an array
-- lets the encode-time JIT trace stay attached. Array fields and map
-- fields aren't recursed into: phase 2 mutation is on MessageView only.
function MessageView:is_dirty()
    if self._eager_only then return self._eager_dirty == true end
    local d = self._dirty
    if d ~= nil and next(d) ~= nil then return true end
    local subs = self._sub_msg_views
    for i = 1, #subs do
        if subs[i]:is_dirty() then return true end
    end
    return false
end

-- Materialize a possibly-view value so encode_field can consume it.
-- Lazy sub-views are converted to plain tables; raw values pass through.
local function materialize(value)
    if type(value) ~= 'table' then return value end
    local m = getmetatable(value)
    if m == MessageView then
        return value:totable()
    elseif m == ArrayView then
        return value:tolist()
    elseif m == MapView then
        return value:totable()
    end
    return value
end

-- Recursively materialize a value to a plain Lua structure suitable for
-- the eager codec. For a MessageView, we walk every present field —
-- this is the cost of mutation; pure-read messages stay lazy.
function MessageView:totable()
    if self._eager_only then return self._eager end
    local out = {}
    for name in self:names() do
        local v = self:get(name)
        out[name] = materialize(v)
    end
    -- Preserve unknown fields for round-trip.
    if self._segments then
        local unknown = {}
        local fbi = self._desc.field_by_id
        for i = 1, #self._segments do
            local seg = self._segments[i]
            if fbi[seg.id] == nil then
                unknown[#unknown + 1] = self._bytes:sub(seg.tag_start, seg.next_start - 1)
            end
        end
        if #unknown > 0 then out._unknown_fields = table.concat(unknown) end
    end
    return out
end

-- :encode() emits bytes. Three modes:
--   1. WKT eager-wrap: delegate to desc.encode on the materialized table.
--   2. Untouched (no dirty fields, no dirty sub-views): emit the original
--      bytes verbatim — perfect byte-for-byte passthrough.
--   3. Mixed: walk fields in id order, splice clean segments, encode dirty
--      values fresh. Unknown segments are emitted at the end.
function MessageView:encode()
    if self._eager_only then
        return self._desc.encode(self._eager)
    end
    if not self:is_dirty() then
        return self._bytes
    end

    local out = {}
    local fields = self._desc.fields
    local bytes = self._bytes
    local by_id = self._by_id
    local dirty = self._dirty or {}
    local cache = self._cache

    -- Active-oneof resolution mirrors codec.encode_message:
    -- a oneof field is only emitted if it's the active branch.
    local active
    local oneofs_list = self._desc.oneofs_list
    if oneofs_list then
        active = {}
        for i = 1, #oneofs_list do
            local oo = oneofs_list[i]
            local members = oo.members
            for j = 1, #members do
                local fname = members[j]
                if dirty[fname] or by_id[self._desc.field_by_name[fname].id] then
                    active[oo.name] = fname
                end
            end
        end
    end

    for i = 1, #fields do
        local f = fields[i]
        local fname = f.name
        local is_dirty = dirty[fname] == true
        local is_sub_dirty = false
        if not is_dirty then
            local cached = cache[fname]
            if cached ~= nil and type(cached) == 'table'
                and getmetatable(cached) == MessageView and cached:is_dirty() then
                is_sub_dirty = true
            end
        end

        if f.oneof and active and active[f.oneof] ~= fname then
            -- Inactive oneof branch: skip entirely.
        elseif is_dirty or is_sub_dirty then
            local v = cache[fname]
            v = materialize(v)
            codec.encode_field(f, v, out, f.optional or (f.oneof ~= nil))
        elseif by_id[f.id] then
            local segs = by_id[f.id]
            for j = 1, #segs do
                local seg = segs[j]
                out[#out + 1] = bytes:sub(seg.tag_start, seg.next_start - 1)
            end
        end
    end

    -- Unknown segments preserved at the end (matches codec's _unknown_fields
    -- trailer convention).
    local fbi = self._desc.field_by_id
    local segments = self._segments
    for i = 1, #segments do
        local seg = segments[i]
        if fbi[seg.id] == nil then
            out[#out + 1] = bytes:sub(seg.tag_start, seg.next_start - 1)
        end
    end

    return table.concat(out)
end

-- ---------------------------------------------------------------------------
-- Public entry
-- ---------------------------------------------------------------------------

M.MessageView = MessageView
M.ArrayView   = ArrayView
M.MapView     = MapView

return M