local TXSE = require("txse-lib")

-- Header structure constants
-- UdpHeader: session(8) + sequence(8) + messageCount(2) + packetType(1) = 19 bytes
local HEADER_SIZE = 19
local MESSAGE_COUNT_OFFSET = 16
local PACKET_TYPE_OFFSET = 18

-- Packet types from RakeUdpPacketType enum
local PACKET_TYPE_SEQUENCED_DATA = 0

-- UdpSequencedMessage structure constants
-- Message format: length(2) + streamId(1) + payload(length-1)
local MSG_LENGTH_FIELD_SIZE = 2
local MSG_HEADER_MIN_SIZE = 3  -- length(2) + streamId(1)

-- Get the generated protocol and parsers
local generated = require("txse-rake-udp-generated")
local proto, messages = generated.register({})

-- Look up the UdpHeader and UdpSequencedMessage parsers once at load time
local header_msg = nil
local sequenced_msg = nil
if messages.messages then
    for _, msg in ipairs(messages.messages) do
        if msg.name == "UdpHeader" then
            header_msg = msg
        elseif msg.name == "UdpSequencedMessage" then
            sequenced_msg = msg
        end
    end
end

if not header_msg then
    error("Failed to find UdpHeader message in generated protocol")
end
if not sequenced_msg then
    error("Failed to find UdpSequencedMessage message in generated protocol")
end

-- Create a dissector function that only parses the UdpHeader
function proto.dissector(buffer, pinfo, tree)
    local length = buffer:len()

    if length < HEADER_SIZE then
        return 0
    end

    pinfo.cols.protocol = proto.name

    -- Parse the UdpHeader
    local subtree = tree:add(proto, buffer(0, HEADER_SIZE), "RAKE UDP Header")
    header_msg:parse(buffer(0, HEADER_SIZE):tvb(), pinfo, subtree)

    -- Extract packet type from header
    local packet_type = buffer(PACKET_TYPE_OFFSET, 1):uint()

    -- Only SEQUENCED_DATA has messages
    if packet_type ~= PACKET_TYPE_SEQUENCED_DATA then
        return HEADER_SIZE
    end

    -- Extract message count from header (little-endian)
    local message_count = buffer(MESSAGE_COUNT_OFFSET, 2):le_int()

    -- Parse each UdpSequencedMessage after the header
    local offset = HEADER_SIZE
    for i = 1, message_count do
        if offset + MSG_HEADER_MIN_SIZE > length then
            break
        end

        local msg_length = buffer(offset, MSG_LENGTH_FIELD_SIZE):le_int()
        local total_msg_size = MSG_LENGTH_FIELD_SIZE + msg_length

        if offset + total_msg_size > length then
            -- Message extends beyond packet
            break
        end

        -- Parse this UdpSequencedMessage
        local msg_buffer = buffer(offset, total_msg_size)
        local msg_tree = subtree:add(proto, msg_buffer, string.format("Message %d", i))
        sequenced_msg:parse(msg_buffer:tvb(), pinfo, msg_tree)

        -- Route the payload to FEED dissector
        -- Payload length is msg_length - 1 (streamId is included in msg_length)
        local payload_length = msg_length - 1
        if payload_length > 0 then
            local payload_buffer = buffer(offset + MSG_HEADER_MIN_SIZE, payload_length)
            TXSE.dissector("txse_feed"):call(payload_buffer:tvb(), pinfo, msg_tree)
        end

        offset = offset + total_msg_size
    end

    return offset
end

-- Heuristic dissector using structural validation
proto:register_heuristic("udp", function(buffer, pinfo, tree)
    if buffer:len() < HEADER_SIZE then
        return false
    end

    -- Session format (8 bytes): 32-bit timestamp (seconds since epoch) + 4 ASCII chars (environment)
    -- In little-endian: bytes 0-3 are timestamp, bytes 4-7 are environment

    -- Extract timestamp (bytes 0-3 as little-endian uint32)
    local timestamp = buffer(0, 4):le_uint()

    -- Check environment bytes (bytes 4-7) are printable ASCII
    for i = 4, 7 do
        if not TXSE.is_printable_ascii(buffer(i, 1):uint()) then
            return false
        end
    end

    -- Check timestamp is reasonable: >= Jan 1, 2025 (1735689600) and < Jan 1, 2035 (2051222400)
    if timestamp < 1735689600 or timestamp >= 2051222400 then
        return false
    end

    -- Check message count is reasonable (0-255)
    local message_count = buffer(MESSAGE_COUNT_OFFSET, 2):le_int()
    if message_count < 0 or message_count > 255 then
        return false
    end

    -- Check packet type is valid (0-5)
    local packet_type = buffer(PACKET_TYPE_OFFSET, 1):uint()
    if packet_type > 5 then
        return false
    end

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

-- Register for manual decode-as usage
DissectorTable.get("udp.port"):add_for_decode_as(proto)
