/**

 * vim: set ts=4 :

 * =============================================================================

 * SourceMod

 * Copyright (C) 2004-2010 AlliedModders LLC.  All rights reserved.

 * =============================================================================

 *

 * This program is free software; you can redistribute it and/or modify it under

 * the terms of the GNU General Public License, version 3.0, as published by the

 * Free Software Foundation.

 * 

 * This program is distributed in the hope that it will be useful, but WITHOUT

 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS

 * FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more

 * details.

 *

 * You should have received a copy of the GNU General Public License along with

 * this program.  If not, see <http://www.gnu.org/licenses/>.

 *

 * As a special exception, AlliedModders LLC gives you permission to link the

 * code of this program (as well as its derivative works) to "Half-Life 2," the

 * "Source Engine," the "SourcePawn JIT," and any Game MODs that run on software

 * by the Valve Corporation.  You must obey the GNU General Public License in

 * all respects for all other code used.  Additionally, AlliedModders LLC grants

 * this exception to all derivative works.  AlliedModders LLC defines further

 * exceptions, found in LICENSE.txt (as of this writing, version JULY-31-2007),

 * or <http://www.sourcemod.net/license.php>.

 *

 * Version: $Id$

 */

#include "UserMessages.h"

#include "sm_stringutil.h"

#if SOURCE_ENGINE == SE_CSGO

#include <cstrike15_usermessage_helpers.h>

#endif

UserMessages g_UserMsgs;

#if SOURCE_ENGINE == SE_CSGO

SH_DECL_HOOK3_void(IVEngineServer, SendUserMessage, SH_NOATTRIB, 0, IRecipientFilter &, int, const protobuf::Message &);

#else

#if SOURCE_ENGINE >= SE_LEFT4DEAD

SH_DECL_HOOK3(IVEngineServer, UserMessageBegin, SH_NOATTRIB, 0, bf_write *, IRecipientFilter *, int, const char *);

#else

SH_DECL_HOOK2(IVEngineServer, UserMessageBegin, SH_NOATTRIB, 0, bf_write *, IRecipientFilter *, int);

#endif

SH_DECL_HOOK0_void(IVEngineServer, MessageEnd, SH_NOATTRIB, 0);

#endif // ==SE_CSGO

UserMessages::UserMessages()

#ifndef USE_PROTOBUF_USERMESSAGES

	: m_InterceptBuffer(m_pBase, 2500)

{

	m_Names = sm_trie_create();

#else

	: m_InterceptBuffer(NULL)

{

#endif

	m_HookCount = 0;

	m_InExec = false;

	m_InHook = false;

	m_CurFlags = 0;

	m_CurId = INVALID_MESSAGE_ID;

}

UserMessages::~UserMessages()

{

#ifndef USE_PROTOBUF_USERMESSAGES

	sm_trie_destroy(m_Names);

#endif

	CStack<ListenerInfo *>::iterator iter;

	for (iter=m_FreeListeners.begin(); iter!=m_FreeListeners.end(); iter++)

	{

		delete (*iter);

	}

	m_FreeListeners.popall();

}

void UserMessages::OnSourceModStartup(bool late)

{

#ifndef USE_PROTOBUF_USERMESSAGES

	/* -1 means SourceMM was unable to get the user message list */

	m_FallbackSearch = (g_SMAPI->GetUserMessageCount() == -1);

#endif

}

void UserMessages::OnSourceModAllInitialized()

{

	g_ShareSys.AddInterface(NULL, this);

}

void UserMessages::OnSourceModAllShutdown()

{

	if (m_HookCount)

	{

#if SOURCE_ENGINE == SE_CSGO

		SH_REMOVE_HOOK_MEMFUNC(IVEngineServer, SendUserMessage, engine, this, &UserMessages::OnSendUserMessage_Pre, false);

		SH_REMOVE_HOOK_MEMFUNC(IVEngineServer, SendUserMessage, engine, this, &UserMessages::OnSendUserMessage_Post, true);

#else

		SH_REMOVE_HOOK_MEMFUNC(IVEngineServer, UserMessageBegin, engine, this, &UserMessages::OnStartMessage_Pre, false);

		SH_REMOVE_HOOK_MEMFUNC(IVEngineServer, UserMessageBegin, engine, this, &UserMessages::OnStartMessage_Post, true);

		SH_REMOVE_HOOK_MEMFUNC(IVEngineServer, MessageEnd, engine, this, &UserMessages::OnMessageEnd_Pre, false);

		SH_REMOVE_HOOK_MEMFUNC(IVEngineServer, MessageEnd, engine, this, &UserMessages::OnMessageEnd_Post, true);

#endif

	}

	m_HookCount = 0;

}

int UserMessages::GetMessageIndex(const char *msg)

{

#if SOURCE_ENGINE == SE_CSGO

	// Can split this per engine and/or game later

	return g_Cstrike15UsermessageHelpers.GetIndex(msg);

#else

	int msgid;

	if (!sm_trie_retrieve(m_Names, msg, reinterpret_cast<void **>(&msgid)))

	{

		if (m_FallbackSearch)

		{

			char msgbuf[64];

			int size;

			msgid = 0;

			while (gamedll->GetUserMessageInfo(msgid, msgbuf, sizeof(msgbuf), size))

			{

				if (strcmp(msgbuf, msg) == 0)

				{

					sm_trie_insert(m_Names, msg, reinterpret_cast<void *>(msgid));

					return msgid;

				}

				msgid++;

			}

		}

		msgid = g_SMAPI->FindUserMessage(msg);

		if (msgid != INVALID_MESSAGE_ID)

		{

			sm_trie_insert(m_Names, msg, reinterpret_cast<void *>(msgid));

		}

	}

	return msgid;

#endif

}

bool UserMessages::GetMessageName(int msgid, char *buffer, size_t maxlength) const

{

#if SOURCE_ENGINE == SE_CSGO

	const char *pszName = g_Cstrike15UsermessageHelpers.GetName(msgid);

	if (!pszName)

		return false;

	strncopy(buffer, pszName, maxlength);

	return true;

#else

	if (m_FallbackSearch)

	{

		int size;

		return gamedll->GetUserMessageInfo(msgid, buffer, maxlength, size);

	}

	const char *msg = g_SMAPI->GetUserMessage(msgid);

	if (msg)

	{

		strncopy(buffer, msg, maxlength);

		return true;

	}

	return false;

#endif

}

bf_write *UserMessages::StartBitBufMessage(int msg_id, const cell_t players[], unsigned int playersNum, int flags)

{

#ifdef USE_PROTOBUF_USERMESSAGES

	return NULL;

#else

	bf_write *buffer;

	if (m_InExec || m_InHook)

	{

		return NULL;

	}

	if (msg_id < 0 || msg_id >= 255)

	{

		return NULL;

	}

	m_CellRecFilter.Initialize(players, playersNum);

	m_CurFlags = flags;

	if (m_CurFlags & USERMSG_INITMSG)

	{

		m_CellRecFilter.SetToInit(true);

	}

	if (m_CurFlags & USERMSG_RELIABLE)

	{

		m_CellRecFilter.SetToReliable(true);

	}

	m_InExec = true;

	if (m_CurFlags & USERMSG_BLOCKHOOKS)

	{

#if SOURCE_ENGINE >= SE_LEFT4DEAD

		buffer = ENGINE_CALL(UserMessageBegin)(static_cast<IRecipientFilter *>(&m_CellRecFilter), msg_id, g_SMAPI->GetUserMessage(msg_id));

#else

		buffer = ENGINE_CALL(UserMessageBegin)(static_cast<IRecipientFilter *>(&m_CellRecFilter), msg_id);

#endif

	} else {

#if SOURCE_ENGINE >= SE_LEFT4DEAD

		buffer = engine->UserMessageBegin(static_cast<IRecipientFilter *>(&m_CellRecFilter), msg_id, g_SMAPI->GetUserMessage(msg_id));

#else

		buffer = engine->UserMessageBegin(static_cast<IRecipientFilter *>(&m_CellRecFilter), msg_id);

#endif

	}

	return buffer;

#endif // USE_PROTOBUF_USERMESSAGES

}

google::protobuf::Message *UserMessages::StartProtobufMessage(int msg_id, const cell_t players[], unsigned int playersNum, int flags)

{

#ifndef USE_PROTOBUF_USERMESSAGES

	return NULL;

#else

	protobuf::Message *buffer;

	if (m_InExec || m_InHook)

	{

		return NULL;

	}

	if (msg_id < 0 || msg_id >= 255)

	{

		return NULL;

	}

	m_CurId = msg_id;

	m_CellRecFilter.Initialize(players, playersNum);

	m_CurFlags = flags;

	if (m_CurFlags & USERMSG_INITMSG)

	{

		m_CellRecFilter.SetToInit(true);

	}

	if (m_CurFlags & USERMSG_RELIABLE)

	{

		m_CellRecFilter.SetToReliable(true);

	}

	m_InExec = true;

	if (m_CurFlags & USERMSG_BLOCKHOOKS)

	{

		// direct message creation, return buffer "from engine". keep track

		m_FakeEngineBuffer = g_Cstrike15UsermessageHelpers.GetPrototype(msg_id)->New();

		buffer = m_FakeEngineBuffer;

	} else {

		protobuf::Message *msg = OnStartMessage_Pre(static_cast<IRecipientFilter *>(&m_CellRecFilter), msg_id, g_SMAPI->GetUserMessage(msg_id));

		switch (m_FakeMetaRes)

		{

		case MRES_IGNORED:

		case MRES_HANDLED:

			m_FakeEngineBuffer = g_Cstrike15UsermessageHelpers.GetPrototype(msg_id)->New();

			buffer = m_FakeEngineBuffer;

			break;		

		case MRES_OVERRIDE:

			m_FakeEngineBuffer = g_Cstrike15UsermessageHelpers.GetPrototype(msg_id)->New();

		// fallthrough

		case MRES_SUPERCEDE:

			buffer = msg;

			break;

		}

		OnStartMessage_Post(static_cast<IRecipientFilter *>(&m_CellRecFilter), msg_id, g_SMAPI->GetUserMessage(msg_id));

	}

	return buffer;

#endif // USE_PROTOBUF_USERMESSAGES

}

bool UserMessages::EndMessage()

{

	if (!m_InExec)

	{

		return false;

	}

#if SOURCE_ENGINE == SE_CSGO

	if (m_CurFlags & USERMSG_BLOCKHOOKS)

	{

		ENGINE_CALL(SendUserMessage)(static_cast<IRecipientFilter &>(m_CellRecFilter), m_CurId, *m_FakeEngineBuffer);

		delete m_FakeEngineBuffer;

		m_FakeEngineBuffer = NULL;

	} else {

		OnMessageEnd_Pre();

		switch (m_FakeMetaRes)

		{

		case MRES_IGNORED:

		case MRES_HANDLED:

		case MRES_OVERRIDE:

			engine->SendUserMessage(static_cast<IRecipientFilter &>(m_CellRecFilter), m_CurId, *m_FakeEngineBuffer);

			delete m_FakeEngineBuffer;

			m_FakeEngineBuffer = NULL;

			break;

		//case MRES_SUPERCEDE:

		}

		OnMessageEnd_Post();

	}

#else

	if (m_CurFlags & USERMSG_BLOCKHOOKS)

	{

		ENGINE_CALL(MessageEnd)();

	} else {

		engine->MessageEnd();

	}

#endif // SE_CSGO

	m_InExec = false;

	m_CurFlags = 0;

	m_CellRecFilter.Reset();

	return true;

}

UserMessageType UserMessages::GetUserMessageType() const

{

#ifdef USE_PROTOBUF_USERMESSAGES

	return UM_Protobuf;

#else

	return UM_BitBuf;

#endif

}

bool UserMessages::HookUserMessage2(int msg_id,

									IUserMessageListener *pListener,

									bool intercept)

{

#ifdef USE_PROTOBUF_USERMESSAGES

	return InternalHook(msg_id, (IProtobufUserMessageListener *)pListener, intercept, true);

#else

	return InternalHook(msg_id, (IBitBufUserMessageListener *)pListener, intercept, true);

#endif

}

bool UserMessages::UnhookUserMessage2(int msg_id,

									  IUserMessageListener *pListener,

									  bool intercept)

{

#ifdef USE_PROTOBUF_USERMESSAGES

	return InternalUnhook(msg_id, (IProtobufUserMessageListener *)pListener, intercept, true);

#else

	return InternalUnhook(msg_id, (IBitBufUserMessageListener *)pListener, intercept, true);

#endif

}

bool UserMessages::HookUserMessage(int msg_id, IUserMessageListener *pListener, bool intercept)

{

#ifdef USE_PROTOBUF_USERMESSAGES

	return InternalHook(msg_id, (IProtobufUserMessageListener *)pListener, intercept, false);

#else

	return InternalHook(msg_id, (IBitBufUserMessageListener *)pListener, intercept, false);

#endif

}

bool UserMessages::UnhookUserMessage(int msg_id, IUserMessageListener *pListener, bool intercept)

{	

#ifdef USE_PROTOBUF_USERMESSAGES

	return InternalUnhook(msg_id, (IProtobufUserMessageListener *)pListener, intercept, false);

#else

	return InternalUnhook(msg_id, (IBitBufUserMessageListener *)pListener, intercept, false);

#endif

}

#ifdef USE_PROTOBUF_USERMESSAGES

bool UserMessages::InternalHook(int msg_id, IProtobufUserMessageListener *pListener, bool intercept, bool isNew)

#else

bool UserMessages::InternalHook(int msg_id, IBitBufUserMessageListener *pListener, bool intercept, bool isNew)

#endif

{

	if (msg_id < 0 || msg_id >= 255)

	{

		return false;

	}

	ListenerInfo *pInfo;

	if (m_FreeListeners.empty())

	{

		pInfo = new ListenerInfo;

	} else {

		pInfo = m_FreeListeners.front();

		m_FreeListeners.pop();

	}

	pInfo->Callback = pListener;

	pInfo->IsHooked = false;

	pInfo->KillMe = false;

	pInfo->IsNew = isNew;

	if (!m_HookCount++)

	{

#if SOURCE_ENGINE == SE_CSGO

		SH_ADD_HOOK_MEMFUNC(IVEngineServer, SendUserMessage, engine, this, &UserMessages::OnSendUserMessage_Pre, false);

		SH_ADD_HOOK_MEMFUNC(IVEngineServer, SendUserMessage, engine, this, &UserMessages::OnSendUserMessage_Post, true);

#else

		SH_ADD_HOOK_MEMFUNC(IVEngineServer, UserMessageBegin, engine, this, &UserMessages::OnStartMessage_Pre, false);

		SH_ADD_HOOK_MEMFUNC(IVEngineServer, UserMessageBegin, engine, this, &UserMessages::OnStartMessage_Post, true);

		SH_ADD_HOOK_MEMFUNC(IVEngineServer, MessageEnd, engine, this, &UserMessages::OnMessageEnd_Pre, false);

		SH_ADD_HOOK_MEMFUNC(IVEngineServer, MessageEnd, engine, this, &UserMessages::OnMessageEnd_Post, true);

#endif

	}

	if (intercept)

	{

		m_msgIntercepts[msg_id].push_back(pInfo);

	} else {

		m_msgHooks[msg_id].push_back(pInfo);

	}

	return true;

}

#ifdef USE_PROTOBUF_USERMESSAGES

bool UserMessages::InternalUnhook(int msg_id, IProtobufUserMessageListener *pListener, bool intercept, bool isNew)

#else

bool UserMessages::InternalUnhook(int msg_id, IBitBufUserMessageListener *pListener, bool intercept, bool isNew)

#endif

{

	MsgList *pList;

	MsgIter iter;

	ListenerInfo *pInfo;

	bool deleted = false;

	if (msg_id < 0 || msg_id >= 255)

	{

		return false;

	}

	pList = (intercept) ? &m_msgIntercepts[msg_id] : &m_msgHooks[msg_id];

	for (iter=pList->begin(); iter!=pList->end(); iter++)

	{

		pInfo = (*iter);

		if (pInfo->Callback == pListener && pInfo->IsNew == isNew)

		{

			if (pInfo->IsHooked)

			{

				pInfo->KillMe = true;

				return true;

			}

			pList->erase(iter);

			deleted = true;

			break;

		}

	}

	if (deleted)

	{

		_DecRefCounter();

	}

	return deleted;

}

void UserMessages::_DecRefCounter()

{

	if (--m_HookCount == 0)

	{

#if SOURCE_ENGINE == SE_CSGO

		SH_REMOVE_HOOK_MEMFUNC(IVEngineServer, SendUserMessage, engine, this, &UserMessages::OnSendUserMessage_Pre, false);

		SH_REMOVE_HOOK_MEMFUNC(IVEngineServer, SendUserMessage, engine, this, &UserMessages::OnSendUserMessage_Post, true);

#else

		SH_REMOVE_HOOK_MEMFUNC(IVEngineServer, UserMessageBegin, engine, this, &UserMessages::OnStartMessage_Pre, false);

		SH_REMOVE_HOOK_MEMFUNC(IVEngineServer, UserMessageBegin, engine, this, &UserMessages::OnStartMessage_Post, true);

		SH_REMOVE_HOOK_MEMFUNC(IVEngineServer, MessageEnd, engine, this, &UserMessages::OnMessageEnd_Pre, false);

		SH_REMOVE_HOOK_MEMFUNC(IVEngineServer, MessageEnd, engine, this, &UserMessages::OnMessageEnd_Post, true);

#endif

	}

}

#ifdef USE_PROTOBUF_USERMESSAGES

// Unlike protobuf's Message::CopyFrom:

// - Does not require to/from message descriptors to be identical.

// - Fields are copied by name, rather than index.

// - String fields are retrieved for copy with GetStringReference, to avoid destucting libprotobuf std::string in sm.

// - Unknown fields in |from| message are not copied to |to| message.

// - Copying Message fields recursively calls this function, rather than Message::MergeFrom.

static void _CopyProtobufMessage(const protobuf::Message &from, protobuf::Message *to)

{

	const protobuf::Descriptor* descriptor = from.GetDescriptor();

	const protobuf::Reflection* from_reflection = from.GetReflection();

	const protobuf::Reflection* to_reflection = to->GetReflection();

	protobuf::UnknownFieldSet toUnknownFields;

	int fromFieldCount = descriptor->field_count();

	for (int i = 0; i < fromFieldCount; i++)

	{

		const protobuf::FieldDescriptor* fromField = descriptor->field(i);

		const protobuf::FieldDescriptor* toField = to->GetDescriptor()->FindFieldByName(fromField->name());

		// This will stripped new fields off of intercepted usermessages.

		// TODO: fix that.

		if (!toField)

		{

			continue;

		}

		if (fromField->is_repeated())

		{

			int count = from_reflection->FieldSize(from, fromField);

			for (int j = 0; j < count; j++)

			{

				switch (fromField->cpp_type())

				{

#define HANDLE_TYPE(CPPTYPE, METHOD)                                           \

				case protobuf::FieldDescriptor::CPPTYPE_##CPPTYPE:             \

					to_reflection->Add##METHOD(to, toField,                    \

					from_reflection->GetRepeated##METHOD(from, fromField, j)); \

					break;

				HANDLE_TYPE(INT32 , Int32 );

				HANDLE_TYPE(INT64 , Int64 );

				HANDLE_TYPE(UINT32, UInt32);

				HANDLE_TYPE(UINT64, UInt64);

				HANDLE_TYPE(FLOAT , Float );

				HANDLE_TYPE(DOUBLE, Double);

				HANDLE_TYPE(BOOL  , Bool  );

				HANDLE_TYPE(ENUM  , Enum  );

#undef HANDLE_TYPE

				case protobuf::FieldDescriptor::CPPTYPE_STRING:

					{

						std::string buffer;

						to_reflection->AddString(to, toField,

							from_reflection->GetRepeatedStringReference(from, fromField, j, &buffer));

					}

					break;

				case protobuf::FieldDescriptor::CPPTYPE_MESSAGE:

					_CopyProtobufMessage(

						from_reflection->GetRepeatedMessage(from, fromField, j),

						to_reflection->AddMessage(to, toField)

						);

					break;

				}

			}

		}

		else

		{

			switch (fromField->cpp_type())

			{

#define HANDLE_TYPE(CPPTYPE, METHOD)                            \

			case protobuf::FieldDescriptor::CPPTYPE_##CPPTYPE:  \

				to_reflection->Set##METHOD(to, toField,         \

				from_reflection->Get##METHOD(from, fromField)); \

				break;

			HANDLE_TYPE(INT32 , Int32 );

			HANDLE_TYPE(INT64 , Int64 );

			HANDLE_TYPE(UINT32, UInt32);

			HANDLE_TYPE(UINT64, UInt64);

			HANDLE_TYPE(FLOAT , Float );

			HANDLE_TYPE(DOUBLE, Double);

			HANDLE_TYPE(BOOL  , Bool  );

			HANDLE_TYPE(ENUM  , Enum  );

#undef HANDLE_TYPE

			case protobuf::FieldDescriptor::CPPTYPE_STRING:

				{

					std::string buffer;

					to_reflection->SetString(to, toField,

						from_reflection->GetStringReference(from, fromField, &buffer));