require "util"

local p = Proto("pcmeta", "PictoChat Meta")

local message_dissector = Dissector.get("pcmsg")

function p.init()
	local dt = DissectorTable.get("dslmp")
	dt:add(GAMEID.PICTOCHAT, p)
end

local MSG_TYPE = {
	MSG_START = 1,
	MSG_FRAG = 2,

	-- There doesn't seem to be an actual difference between type 4 and 5, even
	-- the sequence number is shared between these message types
	STATUS_A = 4,
	STATUS_B = 5,
}

local MSG_TYPE_MAP = {
	[MSG_TYPE.MSG_START] = "msg start",
	[MSG_TYPE.MSG_FRAG] = "msg fragment",
	[MSG_TYPE.STATUS_A] = "status", -- constantly spammed, contains room users' addresses
	[MSG_TYPE.STATUS_B] = "status",
}

p.fields.unknown = ProtoField.uint16("pcmeta.unknown", "Unknown")

local dissect_msg_type = { }

p.fields.msg_start_len = ProtoField.uint8("pcmeta.msg_start.len", "Total length")
p.fields.msg_start_id = ProtoField.uint32("pcmeta.msg_start.id", "ID")
p.fields.msg_start_seq = ProtoField.uint16("pcmeta.msg_start.seq", "Sequence")
p.fields.msg_start_new = ProtoField.bool("pcmeta.msg_start.new", "New")
dissect_msg_type[MSG_TYPE.MSG_START] = function (buffer, pinfo, tree)
	local subtree = tree:add(p, buffer(), string.format("PictoChat Meta Message start: %d bytes", buffer:len()))

	subtree:add_le(p.fields.unknown, buffer(0x00, 2))
	subtree:add_le(p.fields.unknown, buffer(0x02, 2)) -- usu. ffffh
	subtree:add_le(p.fields.msg_start_len, buffer(0x04, 2))
	subtree:add_le(p.fields.unknown, buffer(0x06, 2)) -- usu. 0000h
	subtree:add_le(p.fields.unknown, buffer(0x08, 2))
	subtree:add_le(p.fields.unknown, buffer(0x0a, 2)) -- usu. 0000h
	subtree:add_le(p.fields.msg_start_id, buffer(0x0c, 4))
	local id = buffer(0x0c, 4):le_uint()
	subtree:add_le(p.fields.msg_start_seq, buffer(0x10, 2))
	subtree:add_le(p.fields.new, buffer(0x12, 2))
	local new = buffer(0x12, 2):le_uint() ~= 0

	if not pinfo.visited then
		if new then
			pc_global.mid = id
		end

		pc_global.pid_mid_map[pinfo.number] = pc_global.mid
		pc_global.msg[pc_global.mid] = ByteArray.new()
	end
end

p.fields.body_len = ProtoField.uint8("pcmeta.body.len", "Length")
p.fields.body_end = ProtoField.bool("pcmeta.body.end", "End")
p.fields.body_offset = ProtoField.uint16("pcmeta.body.offset", "Offset")
p.fields.body = ProtoField.bytes("pcmeta.body", "Data")
p.fields.body_seq = ProtoField.uint16("pcmeta.body.seq", "Sequence")
p.fields.body_new = ProtoField.bool("pcmeta.body.new", "New")
dissect_msg_type[MSG_TYPE.MSG_FRAG] = function (buffer, pinfo, tree)
	local subtree = tree:add(p, buffer(), string.format("PictoChat Meta Message Fragment: %d bytes", buffer:len()))

	subtree:add_le(p.fields.unknown, buffer(0x00, 2))
	subtree:add_le(p.fields.body_len, buffer(0x02, 1))
	local body_len = buffer(0x02, 1):le_uint()

	subtree:add_le(p.fields.body_end, buffer(0x03, 1))
	-- This appears to be some kind of offset for indicating where to store the
	-- current frame's data in a larger buffer. Messages sent in multiple parts
	-- increment this value by 160 for each new original (p.fields.original ==
	-- True) message.
	subtree:add_le(p.fields.body_offset, buffer(0x04, 2))
	local body_offset = buffer(0x04, 2):le_uint()
	subtree:add_le(p.fields.unknown, buffer(0x06, 2)) -- usually 0
	buffer = buffer(0x08)

	-- This appears to be the actual message data (the drawing) sent as an
	-- array of 8x8 tiles.
	local body = buffer(0, body_len)
	subtree:add(p.fields.body, body)

	buffer = buffer(body_len)

	subtree:add_le(p.fields.body_seq, buffer(0x00, 2))
	subtree:add_le(p.fields.body_new, buffer(0x02, 2))
	local new = buffer(0x02, 2):le_uint() ~= 0

	if not pinfo.visited then
		pc_global.pid_mid_map[pinfo.number] = pc_global.mid

		-- reassembly
		if pc_global.msg[pc_global.mid]:len() == body_offset and new == true then
			pc_global.msg[pc_global.mid]:append(body:bytes())
		end
	end
end

p.fields.status_id = ProtoField.uint32("pcmeta.status.id", "ID")
p.fields.status_addr = ProtoField.bytes("pcmeta.status.addr", "Address", base.COLON)
p.fields.status_seq = ProtoField.uint16("pcmeta.status.seq", "Sequence")
p.fields.status_new = ProtoField.bool("pcmeta.status.new", "New")
dissect_msg_type[MSG_TYPE.STATUS_A] = function (buffer, pinfo, tree)
	local subtree = tree:add(p, buffer(), string.format("PictoChat Meta Status: %d bytes", buffer:len()))

	subtree:add_le(p.fields.status_id, buffer(0x00, 4))
	local id = buffer(0x00, 4):le_uint()

	buffer = buffer(0x04)

	for _ = 0, 15 do
		add_addr_le(subtree, p.fields.status_addr, buffer(0, 6))
		buffer = buffer(6)
	end

	subtree:add_le(p.fields.status_seq, buffer(0x00, 2))
	subtree:add_le(p.fields.status_new, buffer(0x02, 2))

	pc_global.pid_mid_map[pinfo.number] = id
end
dissect_msg_type[MSG_TYPE.STATUS_B] = dissect_msg_type[MSG_TYPE.STATUS_A]

p.fields.magic = ProtoField.bytes("pcmeta.magic", "Magic")
p.fields.new = ProtoField.bool("pcmeta.new", "New")
p.fields.type = ProtoField.uint8("pcmeta.type", "Content type", base.DEC, MSG_TYPE_MAP)
p.fields.len = ProtoField.uint16("pcmeta.len", "Payload length")
p.fields.data = ProtoField.bytes("pcmeta.data", "Payload")
function p.dissector(buffer, pinfo, tree)
	-- check magic
	if buffer(0x00, 2):uint() ~= GAMEID.PICTOCHAT then return 0 end
	-- check length (avoid errors)
	if buffer:len() < 6 then return 0 end

	pinfo.cols.protocol = p.name
	local buffer_len = math.min(buffer:len(), 0x0a)
	local subtree = tree:add(p, buffer(0, buffer_len), string.format("PictoChat Meta Header: %d bytes", buffer_len))

  subtree:add_le(p.fields.magic, buffer(0x00, 2))
  subtree:add_le(p.fields.new, buffer(0x02, 2))
	subtree:add_le(p.fields.unknown, buffer(0x04, 2))
	local unknown = buffer(0x04, 2):le_uint()
	if unknown == 0 then
		-- no idea what this field means but if it is 0 then this is the last field
		-- in the pictochat body, and we must return in order to prevent lua errors
		pinfo.cols.info = "???"
		return buffer_len
	end

  subtree:add_le(p.fields.type, buffer(0x06, 2))
	local type = buffer(0x06, 2):le_uint()
	local type_str = MSG_TYPE_MAP[type]

	if type_str ~= nil then
		pinfo.cols.info = type_str
	end

	subtree:add_le(p.fields.len, buffer(0x08, 2))
	local payload_length = buffer(0x08, 2):le_uint()
  local payload = buffer(0x0a, payload_length)

	subtree:add(p.fields.data, payload)

	local subdissector = dissect_msg_type[type]
	if subdissector ~= nil then
		subdissector(payload, pinfo, tree)
	end

	-- prefix info field with message ID
	local mid = pc_global.pid_mid_map[pinfo.number] or 0
	pinfo.cols.info = string.format("[%08x] %s", mid, pinfo.cols.info)

	-- add reassembly to current frame if current frame is part of the complete message
	local reassembly = pc_global.msg[mid]
	if reassembly ~= nil then
		local tvb = ByteArray.tvb(reassembly, string.format("msg %08x", mid))
		message_dissector:call(tvb, pinfo, tree)
	end

	return buffer_len
end