-- MIRP dissector for Wireshark 4.6.x (performance-tuned)
-- 上期能源/上期所 深度行情发布平台接口规范 v1.00 (Chapter 6)

local mirp_proto = Proto("mirp", "MIRP (SHFE Incremental Market Data)")

mirp_proto.prefs.udp_ports = Pref.string(
    "UDP ports",
    "21000,25000",
    "Comma-separated UDP ports, e.g. 12345,12346 or 12340-12345"
)
mirp_proto.prefs.use_heuristic = Pref.bool(
    "Heuristic dissection",
    false,
    "Try to detect MIRP on all UDP (slow on large captures; prefer setting UDP ports)"
)
mirp_proto.prefs.track_change = Pref.bool(
    "Track ChangeNo continuity",
    true,
    "Detect missing increments per InstrumentNo"
)
mirp_proto.prefs.track_packet = Pref.bool(
    "Track PacketNo continuity",
    true,
    "Detect missing incremental packets per TopicID"
)
mirp_proto.prefs.mark_ok_in_tree = Pref.bool(
    "Show OK status in tree",
    false,
    "Add tracking fields even when sequence is OK (slower)"
)
mirp_proto.prefs.auto_color_gap = Pref.bool(
    "Auto-install gap color rules",
    true,
    "Highlight gap packets with black background (like TCP lost segment)"
)

local ef_packet_gap = ProtoExpert.new(
    "mirp.expert.packet_gap",
    "PacketNo gap after packet loss",
    expert.group.SEQUENCE,
    expert.severity.ERROR
)
local ef_change_gap = ProtoExpert.new(
    "mirp.expert.change_gap",
    "ChangeNo gap after missed update",
    expert.group.SEQUENCE,
    expert.severity.ERROR
)
mirp_proto.experts = { ef_packet_gap, ef_change_gap }

local current_udp_ports = {}
local heuristic_registered = false

-- Runtime pref cache (avoid table lookup in hot path)
local cfg = {
    track_change = true,
    track_packet = true,
    mark_ok = false,
}

-- Nested: [topic_id][center_change_no][instrument_no]
local last_change_by_inst = {}
-- Nested: [topic_id][center_change_no]
local last_packet_by_topic = {}

local f_flag          = ProtoField.uint8("mirp.flag", "Flag", base.HEX)
local f_type_id       = ProtoField.uint8("mirp.type_id", "TypeID", base.HEX)
local f_length        = ProtoField.uint16("mirp.length", "Body Length", base.DEC)
local f_packet_no     = ProtoField.int32("mirp.packet_no", "PacketNo", base.DEC)
local f_topic_id      = ProtoField.int16("mirp.topic_id", "TopicID", base.DEC)
local f_snap_millisec = ProtoField.uint16("mirp.snap_millisec", "SnapMillisec", base.DEC)
local f_snap_no       = ProtoField.int32("mirp.snap_no", "SnapNo", base.DEC)
local f_snap_time     = ProtoField.uint32("mirp.snap_time", "SnapTime", base.DEC)
local f_comm_phase_no = ProtoField.uint16("mirp.comm_phase_no", "CommPhaseNo", base.DEC)
local f_center_chg_no = ProtoField.int8("mirp.center_change_no", "CenterChangeNo", base.DEC)
local f_reserved      = ProtoField.int8("mirp.reserved", "Reserved", base.DEC)
local f_field_id      = ProtoField.uint16("mirp.field.id", "FieldID", base.HEX)
local f_field_size    = ProtoField.uint16("mirp.field.size", "FieldSize", base.DEC)
local f_inst_no       = ProtoField.int32("mirp.contract.instrument_no", "InstrumentNo", base.DEC)
local f_change_no     = ProtoField.int32("mirp.contract.change_no", "ChangeNo", base.DEC)
local f_change_expect = ProtoField.int32("mirp.contract.change_expected", "Expected ChangeNo", base.DEC)
local f_change_gap    = ProtoField.uint32("mirp.contract.change_gap", "Missed ChangeNo Count", base.DEC)
local f_packet_expect = ProtoField.int32("mirp.packet_expected", "Expected PacketNo", base.DEC)
local f_packet_gap    = ProtoField.uint32("mirp.packet_gap", "Missed PacketNo Count", base.DEC)
local f_packet_gap_detected = ProtoField.bool("mirp.packet_gap_detected", "Packet Gap Detected")
local f_change_gap_detected = ProtoField.bool("mirp.contract.change_gap_detected", "Change Gap Detected")
local f_gap_loss          = ProtoField.uint8("mirp.gap_loss", "MIRP Sequence Gap", base.DEC)

local CHANGE_STATUS = {
    [0] = "First seen",
    [1] = "OK (continuous)",
    [2] = "Gap (packet loss)",
    [3] = "Duplicate",
    [4] = "Rewind (out-of-order or reset)",
}
local f_change_status = ProtoField.uint8("mirp.contract.change_status", "ChangeNo Status", base.DEC, CHANGE_STATUS)
local f_packet_status = ProtoField.uint8("mirp.packet_status", "PacketNo Status", base.DEC, CHANGE_STATUS)

mirp_proto.fields = {
    f_flag, f_type_id, f_length, f_packet_no, f_topic_id,
    f_snap_millisec, f_snap_no, f_snap_time, f_comm_phase_no,
    f_center_chg_no, f_reserved, f_field_id, f_field_size,
    f_inst_no, f_change_no, f_change_expect, f_change_gap, f_change_status,
    f_packet_expect, f_packet_gap, f_packet_status,
    f_packet_gap_detected, f_change_gap_detected, f_gap_loss,
}

local MIRP_HEADER_LEN = 24
local MIRP_MAX_PACKET = 1232
local FIELD_CONTRACT_HDR = 0x0003

local TYPE_NAMES = {
    [0x00] = "Heartbeat",
    [0x01] = "Incremental Refresh",
}

local function refresh_cfg()
    cfg.track_change = mirp_proto.prefs.track_change
    cfg.track_packet = mirp_proto.prefs.track_packet
    cfg.mark_ok = mirp_proto.prefs.mark_ok_in_tree
end

local function decode_varint(tvb, offset, pkt_len)
    local zigzag = 0
    local shift = 0
    local pos = offset

    while pos < pkt_len do
        local b = tvb(pos, 1):uint()
        pos = pos + 1
        zigzag = zigzag + bit.lshift(bit.band(b, 0x7F), shift)
        if bit.band(b, 0x80) == 0 then
            return zigzag, pos - offset
        end
        shift = shift + 7
        if shift > 63 then
            return nil, pos - offset
        end
    end
    return nil, pos - offset
end

local function decode_vint(tvb, offset, pkt_len)
    local zigzag, consumed = decode_varint(tvb, offset, pkt_len)
    if zigzag == nil then
        return nil, consumed
    end
    return bit.bxor(bit.rshift(zigzag, 1), -bit.band(zigzag, 1)), consumed
end

local function read_vint(tvb, offset, pkt_len)
    return decode_vint(tvb, offset, pkt_len)
end

local function mark_gap_loss(tree, track, gap_flag_pf, expert_pf)
    if not track or track.status_code ~= 2 then
        return
    end
    tree:add(gap_flag_pf, true)
    if expert_pf and track.message then
        tree:add_proto_expert_info(expert_pf, track.message)
    end
end

local function add_seq_fields(tree, track, status_pf, expect_pf, gap_pf, gap_flag_pf, expert_pf)
    if not track then
        return
    end
    if track.status_code == 1 and not cfg.mark_ok then
        return
    end
    tree:add(status_pf, track.status_code)
    if track.status_code ~= 0 and track.status_code ~= 1 then
        tree:add(expect_pf, track.expected)
    end
    if track.gap > 0 then
        tree:add(gap_pf, track.gap)
    end
    mark_gap_loss(tree, track, gap_flag_pf, expert_pf)
    if track.message and track.status_code ~= 2 then
        if track.status_code == 3 or track.status_code == 4 then
            tree:add_expert_info(PI_PROTOCOL, PI_WARN, track.message)
        end
    end
end

local function get_packet_prev(topic_id, center_change_no)
    local t = last_packet_by_topic[topic_id]
    return t and t[center_change_no] or nil
end

local function set_packet_prev(topic_id, center_change_no, packet_no, frame_num)
    local t = last_packet_by_topic[topic_id]
    if not t then
        t = {}
        last_packet_by_topic[topic_id] = t
    end
    t[center_change_no] = { packet_no = packet_no, frame_num = frame_num }
end

local function check_packet_continuity(topic_id, center_change_no, packet_no, type_id, frame_num)
    if type_id ~= 0x01 then
        return nil
    end

    local prev = get_packet_prev(topic_id, center_change_no)
    local track = { status_code = 0, gap = 0, expected = packet_no, message = nil }

    if prev then
        local expected = prev.packet_no + 1
        track.expected = expected
        if packet_no == expected then
            track.status_code = 1
        elseif packet_no > expected then
            track.status_code = 2
            track.gap = packet_no - expected
            track.message = string.format(
                "TopicID %d: PacketNo gap, expected %d got %d (missed %d, last frame %d)",
                topic_id, expected, packet_no, track.gap, prev.frame_num)
        elseif packet_no == prev.packet_no then
            track.status_code = 3
            track.message = string.format(
                "TopicID %d: duplicate PacketNo %d (last frame %d)",
                topic_id, packet_no, prev.frame_num)
        else
            track.status_code = 4
            track.message = string.format(
                "TopicID %d: PacketNo rewind %d -> %d (last frame %d)",
                topic_id, prev.packet_no, packet_no, prev.frame_num)
        end
    end

    set_packet_prev(topic_id, center_change_no, packet_no, frame_num)
    return track
end

local function get_change_prev(topic_id, center_change_no, instrument_no)
    local t = last_change_by_inst[topic_id]
    if not t then return nil end
    local c = t[center_change_no]
    if not c then return nil end
    return c[instrument_no]
end

local function set_change_prev(topic_id, center_change_no, instrument_no, change_no, packet_no, frame_num)
    local t = last_change_by_inst[topic_id]
    if not t then
        t = {}
        last_change_by_inst[topic_id] = t
    end
    local c = t[center_change_no]
    if not c then
        c = {}
        t[center_change_no] = c
    end
    c[instrument_no] = { change_no = change_no, packet_no = packet_no, frame_num = frame_num }
end

local function check_change_continuity(topic_id, center_change_no, instrument_no, change_no, packet_no, frame_num)
    local prev = get_change_prev(topic_id, center_change_no, instrument_no)
    local track = { status_code = 0, gap = 0, expected = change_no, message = nil }

    if prev then
        local expected = prev.change_no + 1
        track.expected = expected
        if change_no == expected then
            track.status_code = 1
        elseif change_no > expected then
            track.status_code = 2
            track.gap = change_no - expected
            track.message = string.format(
                "InstrumentNo %d: ChangeNo gap, expected %d got %d (missed %d, last frame %d)",
                instrument_no, expected, change_no, track.gap, prev.frame_num)
        elseif change_no == prev.change_no then
            track.status_code = 3
            track.message = string.format(
                "InstrumentNo %d: duplicate ChangeNo %d (last frame %d)",
                instrument_no, change_no, prev.frame_num)
        else
            track.status_code = 4
            track.message = string.format(
                "InstrumentNo %d: ChangeNo rewind %d -> %d (last frame %d)",
                instrument_no, prev.change_no, change_no, prev.frame_num)
        end
    end

    set_change_prev(topic_id, center_change_no, instrument_no, change_no, packet_no, frame_num)
    return track
end

local function dissect_contract_header(field_tree, tvb, body_offset, body_len, ctx, pkt_len)
    local pos = body_offset
    local end_pos = body_offset + body_len
    if pos >= end_pos then
        return nil, nil, nil
    end

    local inst_no, c1 = read_vint(tvb, pos, pkt_len)
    if inst_no == nil then
        return nil, nil, nil
    end
    field_tree:add(f_inst_no, tvb(pos, c1))
    pos = pos + c1

    if pos >= end_pos then
        return inst_no, nil, nil
    end

    local change_no, c2 = read_vint(tvb, pos, pkt_len)
    if change_no == nil then
        return inst_no, nil, nil
    end
    local change_tvb = tvb(pos, c2)
    field_tree:add(f_change_no, change_tvb)

    local track = nil
    if cfg.track_change then
        track = check_change_continuity(ctx.topic_id, ctx.center_change_no, inst_no, change_no, ctx.packet_no, ctx.frame_num)
        add_seq_fields(field_tree, track, f_change_status, f_change_expect, f_change_gap,
            f_change_gap_detected, ef_change_gap)
    end
    return inst_no, change_no, track
end

local function dissect_body(tvb, tree, offset, body_len, ctx, pkt_len)
    if body_len <= 0 then
        return
    end

    local body_tree = tree:add(tvb(offset, body_len), "MIRP Body")
    local end_offset = offset + body_len
    local pos = offset
    local contract_idx = 0

    while pos + 4 <= end_offset do
        local field_id = tvb(pos, 2):le_uint()
        local field_size = tvb(pos + 2, 2):le_uint()
        local body_start = pos + 4
        local body_end = body_start + field_size

        if body_end > end_offset then
            break
        end

        if field_id == FIELD_CONTRACT_HDR then
            contract_idx = contract_idx + 1
            local field_tree = body_tree:add(tvb(pos, 4 + field_size),
                string.format("Contract Header #%d", contract_idx))
            field_tree:add_le(f_field_id, tvb(pos, 2))
            field_tree:add_le(f_field_size, tvb(pos + 2, 2))
            local inst_no, change_no, track = dissect_contract_header(
                field_tree, tvb, body_start, field_size, ctx, pkt_len)
            if inst_no ~= nil then
                ctx.contracts[#ctx.contracts + 1] = {
                    instrument_no = inst_no,
                    change_no = change_no,
                    track = track,
                }
            end
        end

        pos = body_end
    end
end

function mirp_proto.dissector(tvb, pinfo, tree)
    local pkt_len = tvb:len()
    if pkt_len < MIRP_HEADER_LEN then
        return
    end

    pinfo.cols.protocol = "MIRP"

    local flag = tvb(0, 1):uint()
    local type_id = tvb(1, 1):uint()
    local body_len = tvb(2, 2):le_uint()
    local packet_no = tvb(4, 4):le_int()
    local topic_id = tvb(8, 2):le_int()
    local center_change_no = tvb(22, 1):int()

    local subtree = tree:add(tvb(), "MIRP Protocol")
    local hdr_tree = subtree:add(tvb(0, MIRP_HEADER_LEN), "MIRP Header")

    hdr_tree:add(f_flag, tvb(0, 1))
    hdr_tree:add(f_type_id, tvb(1, 1))
    hdr_tree:add_le(f_length, tvb(2, 2))
    hdr_tree:add_le(f_packet_no, tvb(4, 4))
    hdr_tree:add_le(f_topic_id, tvb(8, 2))
    hdr_tree:add_le(f_snap_millisec, tvb(10, 2))
    hdr_tree:add_le(f_snap_no, tvb(12, 4))
    hdr_tree:add_le(f_snap_time, tvb(16, 4))
    hdr_tree:add_le(f_comm_phase_no, tvb(20, 2))
    hdr_tree:add(f_center_chg_no, tvb(22, 1))
    hdr_tree:add(f_reserved, tvb(23, 1))

    local packet_track = nil
    if cfg.track_packet then
        packet_track = check_packet_continuity(topic_id, center_change_no, packet_no, type_id, pinfo.number)
        add_seq_fields(hdr_tree, packet_track, f_packet_status, f_packet_expect, f_packet_gap,
            f_packet_gap_detected, ef_packet_gap)
    end

    local ctx = {
        topic_id = topic_id,
        center_change_no = center_change_no,
        packet_no = packet_no,
        frame_num = pinfo.number,
        contracts = {},
    }

    if body_len > 0 and MIRP_HEADER_LEN + body_len <= pkt_len then
        dissect_body(tvb, subtree, MIRP_HEADER_LEN, body_len, ctx, pkt_len)
    end

    local type_name = TYPE_NAMES[type_id] or "Unknown"
    local info = string.format("TopicID=%d PacketNo=%d %s", topic_id, packet_no, type_name)

    if packet_track and packet_track.status_code == 2 then
        info = info .. string.format(" **PACKET GAP+%d**", packet_track.gap)
    elseif packet_track and packet_track.status_code == 4 then
        info = info .. " **PACKET REWIND**"
    end

    local n = #ctx.contracts
    local has_change_gap = false
    if n > 0 then
        local parts = {}
        for i = 1, n do
            local c = ctx.contracts[i]
            local tr = c.track
            if tr and tr.status_code == 2 then
                has_change_gap = true
                parts[i] = string.format("InstNo=%d ChgNo=%d GAP+%d", c.instrument_no, c.change_no, tr.gap)
            elseif c.change_no ~= nil then
                parts[i] = string.format("InstNo=%d ChgNo=%d", c.instrument_no, c.change_no)
            else
                parts[i] = string.format("InstNo=%d", c.instrument_no)
            end
        end
        if has_change_gap then
            info = info .. " **CHANGE GAP**"
        end
        info = info .. " [" .. table.concat(parts, ", ") .. "]"
    end

    -- Packet-level marker for coloring rules (bool sub-fields may not match reliably)
    if (packet_track and packet_track.status_code == 2) or has_change_gap then
        subtree:add(f_gap_loss, 1)
    end

    pinfo.cols.info = info
end

local function heuristic_mirp(tvb, pinfo, tree)
    if tvb:len() < MIRP_HEADER_LEN then
        return false
    end
    local b0 = tvb(0, 1):uint()
    if bit.band(b0, 0x0F) ~= 1 then
        return false
    end
    local type_id = tvb(1, 1):uint()
    if type_id > 1 then
        return false
    end
    local body_len = tvb(2, 2):le_uint()
    local total_len = MIRP_HEADER_LEN + body_len
    if total_len > MIRP_MAX_PACKET or total_len > tvb:len() then
        return false
    end
    if type_id == 0x00 and body_len ~= 0 then
        return false
    end
    mirp_proto.dissector(tvb, pinfo, tree)
    return true
end

local function parse_ports(spec)
    local ports = {}
    local seen = {}
    if not spec or spec == "" then
        return ports
    end

    for part in string.gmatch(spec, "[^,%s]+") do
        local from_s, to_s = part:match("^(%d+)%-(%d+)$")
        if from_s and to_s then
            local from_p = tonumber(from_s)
            local to_p = tonumber(to_s)
            if from_p and to_p and from_p <= to_p then
                for p = from_p, to_p do
                    if p > 0 and p <= 65535 and not seen[p] then
                        seen[p] = true
                        ports[#ports + 1] = p
                    end
                end
            end
        else
            local p = tonumber(part)
            if p and p > 0 and p <= 65535 and not seen[p] then
                seen[p] = true
                ports[#ports + 1] = p
            end
        end
    end

    table.sort(ports)
    return ports
end

local function register_port(port)
    if port and port > 0 and port <= 65535 then
        DissectorTable.get("udp.port"):add(port, mirp_proto)
    end
end

local function unregister_port(port)
    if port and port > 0 and port <= 65535 then
        DissectorTable.get("udp.port"):remove(port, mirp_proto)
    end
end

local function register_ports(ports)
    for i = 1, #ports do
        register_port(ports[i])
    end
end

local function unregister_ports(ports)
    for i = 1, #ports do
        unregister_port(ports[i])
    end
end

local function ensure_heuristic()
    if mirp_proto.prefs.use_heuristic and not heuristic_registered then
        mirp_proto:register_heuristic("udp", heuristic_mirp)
        heuristic_registered = true
    end
end

-- Wireshark 4.x colorfilters: @name@filter@[bg R,G,B][fg R,G,B] (16-bit RGB)
local COLOR_HEADER = '# DO NOT EDIT THIS FILE!  It was created by Wireshark\n'
local COLOR_RULE = '@MIRP Gap Loss@mirp.gap_loss == 1@[0,0,0][65535,65535,65535]\n'

local function strip_managed_color_lines(content)
    if content == "" then
        return ""
    end
    local kept = {}
    for line in content:gmatch("[^\r\n]+") do
        if not line:find("^# DO NOT EDIT") and not line:find("@MIRP ", 1, true) then
            kept[#kept + 1] = line
        end
    end
    if #kept == 0 then
        return ""
    end
    return table.concat(kept, "\n") .. "\n"
end

local function read_global_colorfilters()
    local candidates = {}
    if Dir and Dir.global_config_path then
        candidates[#candidates + 1] = Dir.global_config_path() .. "\\colorfilters"
    end
    candidates[#candidates + 1] = "C:\\Program Files\\Wireshark\\colorfilters"
    for _, candidate in ipairs(candidates) do
        local rf = io.open(candidate, "r")
        if rf then
            local content = rf:read("*a") or ""
            rf:close()
            if content ~= "" then
                return content
            end
        end
    end
    return ""
end

local function install_gap_color_rules()
    if not mirp_proto.prefs.auto_color_gap then
        return
    end
    if not Dir or not Dir.personal_config_path then
        return
    end

    local path = Dir.personal_config_path() .. "\\colorfilters"
    local existing = ""
    local rf = io.open(path, "r")
    if rf then
        existing = rf:read("*a") or ""
        rf:close()
    end

    local rest = strip_managed_color_lines(existing)
    if rest == "" then
        rest = strip_managed_color_lines(read_global_colorfilters())
    end

    local wf = io.open(path, "w")
    if not wf then
        return
    end
    wf:write(COLOR_HEADER)
    wf:write(COLOR_RULE)
    wf:write(rest)
    wf:close()
end

function mirp_proto.init()
    last_change_by_inst = {}
    last_packet_by_topic = {}
    refresh_cfg()
    current_udp_ports = parse_ports(mirp_proto.prefs.udp_ports)
    register_ports(current_udp_ports)
    ensure_heuristic()
    install_gap_color_rules()
end

function mirp_proto.prefs_changed()
    unregister_ports(current_udp_ports)
    current_udp_ports = parse_ports(mirp_proto.prefs.udp_ports)
    register_ports(current_udp_ports)
    refresh_cfg()
    ensure_heuristic()
end