local TXSE = require("txse-lib")
local txse_rake_framing = Proto.new("txse_rake_framing", "TXSE RAKE TCP Framing")

local f_length = ProtoField.uint16("txse_rake_framing.length", "Length", base.DEC)
txse_rake_framing.fields = {f_length}

function txse_rake_framing.dissector(buffer, pinfo, tree)
    local length = buffer:len()
    if length < 2 then
        -- Don't even have a two-byte length? Doesn't look like RAKE.
        return 0
    end

    local offset = 0
    while length - offset > 2 do
        local length_buffer = buffer(offset, 2)
        local msg_length = length_buffer:le_uint()

        local bytes_available_now = length - offset
        local bytes_needed_now = msg_length + 2
        if bytes_available_now < bytes_needed_now then
            -- TCP segment stitching into the next segment and we know exactly how many bytes required
            pinfo.desegment_offset = offset
            pinfo.desegment_len = bytes_needed_now - bytes_available_now
            return length
        end

        local frame_buffer = buffer(offset, bytes_needed_now)
        local payload_buffer = buffer(offset + 2, msg_length)

        local subtree = tree:add(txse_rake_framing, frame_buffer)
        subtree:add_le(f_length, length_buffer)

        TXSE.dissector("txse_rake_tcp"):call(payload_buffer:tvb(), pinfo, subtree)

        offset = offset + bytes_needed_now
    end

    if offset == length then
        -- Consumed the entire buffer
        return length
    elseif length - offset == 1 then
        -- Need at least one more byte to get the length, but cannot tell how many more bytes yet
        pinfo.desegment_offset = offset
        pinfo.desegment_len = DESEGMENT_ONE_MORE_SEGMENT
    elseif length - offset == 2 then
        -- Have exactly 2 bytes (the length field), but need the payload
        local length_buffer = buffer(offset, 2)
        local msg_length = length_buffer:le_uint()
        pinfo.desegment_offset = offset
        pinfo.desegment_len = msg_length
    else
        assert(false)
        return 0
    end
end

txse_rake_framing:register_heuristic("tcp", function(buffer, pinfo, tree)
    -- Must be long enough to have the two-byte length and at least one byte of payload
    if buffer:len() < 3 then
        return false
    end

    local msg_len = buffer(0, 2):le_uint()
    if msg_len < 1 or msg_len > 1500 then
        -- Doesn't seem to be a "reasonable" size embedded in the message
        return false
    end

    -- Check if the payload looks like a valid RAKE message type
    local msg_type = buffer(2, 1):uint()
    -- Valid RAKE TCP message types: '0'-'7' (48-55)
    if msg_type < 48 or msg_type > 55 then
        return false
    end

    return txse_rake_framing.dissector(buffer, pinfo, tree) ~= 0
end)

DissectorTable.get("tcp.port"):add_for_decode_as(txse_rake_framing)
