#include "PacketHandler.h"

#include "BitStream.h"
#include "DluAssert.h"
#include <map>
#include <functional>
#include <cstdio>

void RaknetPacket::Serialize(RakNet::BitStream& bitStream) const {
	bitStream.Write<uint8_t>(static_cast<uint8_t>(messageID));
}

bool RaknetPacket::Deserialize(RakNet::BitStream& bitStream) {
	uint8_t id;
	if (!bitStream.Read<uint8_t>(id)) {
		return false;
	}
	messageID = static_cast<DefaultMessageIDTypes>(id);
	return true;
}

namespace PacketHandler {
	std::map<DefaultMessageIDTypes, std::function<std::unique_ptr<RaknetPacket>()>> g_Handlers = {
		{ID_INTERNAL_PING, []() { return std::make_unique<InternalPing>(); }},
		{ID_PING, []() { return std::make_unique<Ping>(); }},
		{ID_PING_OPEN_CONNECTIONS, []() { return std::make_unique<PingOpenConnections>(); }},
		{ID_CONNECTED_PONG, []() { return std::make_unique<ConnectedPong>(); }},
		{ID_CONNECTION_REQUEST, []() { return std::make_unique<ConnectionRequest>(); }},
		{ID_SECURED_CONNECTION_RESPONSE, []() { return std::make_unique<SecuredConnectionResponse>(); }},
		{ID_SECURED_CONNECTION_CONFIRMATION, []() { return std::make_unique<SecuredConnectionConfirmation>(); }},
		{ID_RPC_MAPPING, []() { return std::make_unique<RPCMapping>(); }},
		{ID_DETECT_LOST_CONNECTIONS, []() { return std::make_unique<DetectLostConnections>(); }},
		{ID_OPEN_CONNECTION_REQUEST, []() { return std::make_unique<OpenConnectionRequest>(); }},
		{ID_OPEN_CONNECTION_REPLY, []() { return std::make_unique<OpenConnectionReply>(); }},
		{ID_RPC, []() { return std::make_unique<RPC>(); }},
		{ID_RPC_REPLY, []() { return std::make_unique<RPCReply>(); }},
		{ID_OUT_OF_BAND_INTERNAL, []() { return std::make_unique<OutOfBandInternal>(); }},
		{ID_CONNECTION_REQUEST_ACCEPTED, []() { return std::make_unique<ConnectionRequestAccepted>(); }},
		{ID_CONNECTION_ATTEMPT_FAILED, []() { return std::make_unique<ConnectionAttemptFailed>(); }},
		{ID_ALREADY_CONNECTED, []() { return std::make_unique<AlreadyConnected>(); }},
		{ID_NEW_INCOMING_CONNECTION, []() { return std::make_unique<NewIncomingConnection>(); }},
		{ID_NO_FREE_INCOMING_CONNECTIONS, []() { return std::make_unique<NoFreeIncomingConnections>(); }},
		{ID_DISCONNECTION_NOTIFICATION, []() { return std::make_unique<DisconnectionNotification>(); }},
		{ID_CONNECTION_LOST, []() { return std::make_unique<ConnectionLost>(); }},
		{ID_RSA_PUBLIC_KEY_MISMATCH, []() { return std::make_unique<RSAPublicKeyMismatch>(); }},
		{ID_CONNECTION_BANNED, []() { return std::make_unique<ConnectionBanned>(); }},
		{ID_INVALID_PASSWORD, []() { return std::make_unique<InvalidPassword>(); }},
		{ID_MODIFIED_PACKET, []() { return std::make_unique<ModifiedPacket>(); }},
		{ID_TIMESTAMP, []() { return std::make_unique<Timestamp>(); }},
		{ID_PONG, []() { return std::make_unique<Pong>(); }},
		{ID_ADVERTISE_SYSTEM, []() { return std::make_unique<AdvertiseSystem>(); }},
		{ID_REMOTE_DISCONNECTION_NOTIFICATION, []() { return std::make_unique<RemoteDisconnectionNotification>(); }},
		{ID_REMOTE_CONNECTION_LOST, []() { return std::make_unique<RemoteConnectionLost>(); }},
		{ID_REMOTE_NEW_INCOMING_CONNECTION, []() { return std::make_unique<RemoteNewIncomingConnection>(); }},
		{ID_DOWNLOAD_PROGRESS, []() { return std::make_unique<DownloadProgress>(); }},
		{ID_FILE_LIST_TRANSFER_HEADER, []() { return std::make_unique<FileListTransferHeader>(); }},
		{ID_FILE_LIST_TRANSFER_FILE, []() { return std::make_unique<FileListTransferFile>(); }},
		{ID_DDT_DOWNLOAD_REQUEST, []() { return std::make_unique<DDTDownloadRequest>(); }},
		{ID_TRANSPORT_STRING, []() { return std::make_unique<TransportString>(); }},
		{ID_REPLICA_MANAGER_CONSTRUCTION, []() { return std::make_unique<ReplicaManagerConstruction>(); }},
		{ID_REPLICA_MANAGER_DESTRUCTION, []() { return std::make_unique<ReplicaManagerDestruction>(); }},
		{ID_REPLICA_MANAGER_SCOPE_CHANGE, []() { return std::make_unique<ReplicaManagerScopeChange>(); }},
		{ID_REPLICA_MANAGER_SERIALIZE, []() { return std::make_unique<ReplicaManagerSerialize>(); }},
		{ID_REPLICA_MANAGER_DOWNLOAD_STARTED, []() { return std::make_unique<ReplicaManagerDownloadStarted>(); }},
		{ID_REPLICA_MANAGER_DOWNLOAD_COMPLETE, []() { return std::make_unique<ReplicaManagerDownloadComplete>(); }},
		{ID_CONNECTION_GRAPH_REQUEST, []() { return std::make_unique<ConnectionGraphRequest>(); }},
		{ID_CONNECTION_GRAPH_REPLY, []() { return std::make_unique<ConnectionGraphReply>(); }},
		{ID_CONNECTION_GRAPH_UPDATE, []() { return std::make_unique<ConnectionGraphUpdate>(); }},
		{ID_CONNECTION_GRAPH_NEW_CONNECTION, []() { return std::make_unique<ConnectionGraphNewConnection>(); }},
		{ID_CONNECTION_GRAPH_CONNECTION_LOST, []() { return std::make_unique<ConnectionGraphConnectionLost>(); }},
		{ID_CONNECTION_GRAPH_DISCONNECTION_NOTIFICATION, []() { return std::make_unique<ConnectionGraphDisconnectionNotification>(); }},
		{ID_ROUTE_AND_MULTICAST, []() { return std::make_unique<RouteAndMulticast>(); }},
		{ID_RAKVOICE_OPEN_CHANNEL_REQUEST, []() { return std::make_unique<RakVoiceOpenChannelRequest>(); }},
		{ID_RAKVOICE_OPEN_CHANNEL_REPLY, []() { return std::make_unique<RakVoiceOpenChannelReply>(); }},
		{ID_RAKVOICE_CLOSE_CHANNEL, []() { return std::make_unique<RakVoiceCloseChannel>(); }},
		{ID_RAKVOICE_DATA, []() { return std::make_unique<RakVoiceData>(); }},
		{ID_AUTOPATCHER_GET_CHANGELIST_SINCE_DATE, []() { return std::make_unique<AutopatcherGetChangelistSinceDate>(); }},
		{ID_AUTOPATCHER_CREATION_LIST, []() { return std::make_unique<AutopatcherCreationList>(); }},
		{ID_AUTOPATCHER_DELETION_LIST, []() { return std::make_unique<AutopatcherDeletionList>(); }},
		{ID_AUTOPATCHER_GET_PATCH, []() { return std::make_unique<AutopatcherGetPatch>(); }},
		{ID_AUTOPATCHER_PATCH_LIST, []() { return std::make_unique<AutopatcherPatchList>(); }},
		{ID_AUTOPATCHER_REPOSITORY_FATAL_ERROR, []() { return std::make_unique<AutopatcherRepositoryFatalError>(); }},
		{ID_AUTOPATCHER_FINISHED_INTERNAL, []() { return std::make_unique<AutopatcherFinishedInternal>(); }},
		{ID_AUTOPATCHER_FINISHED, []() { return std::make_unique<AutopatcherFinished>(); }},
		{ID_AUTOPATCHER_RESTART_APPLICATION, []() { return std::make_unique<AutopatcherRestartApplication>(); }},
		{ID_NAT_PUNCHTHROUGH_REQUEST, []() { return std::make_unique<NATpunchtroughRequest>(); }},
		{ID_NAT_TARGET_NOT_CONNECTED, []() { return std::make_unique<NATTargetNotConnected>(); }},
		{ID_NAT_TARGET_CONNECTION_LOST, []() { return std::make_unique<NATTargetConnectionLost>(); }},
		{ID_NAT_CONNECT_AT_TIME, []() { return std::make_unique<NATConnectAtTime>(); }},
		{ID_NAT_SEND_OFFLINE_MESSAGE_AT_TIME, []() { return std::make_unique<NATSendOfflineMessageAtTime>(); }},
		{ID_NAT_IN_PROGRESS, []() { return std::make_unique<NATInProgress>(); }},
		{ID_DATABASE_QUERY_REQUEST, []() { return std::make_unique<DatabaseQueryRequest>(); }},
		{ID_DATABASE_UPDATE_ROW, []() { return std::make_unique<DatabaseUpdateRow>(); }},
		{ID_DATABASE_REMOVE_ROW, []() { return std::make_unique<DatabaseRemoveRow>(); }},
		{ID_DATABASE_QUERY_REPLY, []() { return std::make_unique<DatabaseQueryReply>(); }},
		{ID_DATABASE_UNKNOWN_TABLE, []() { return std::make_unique<DatabaseUnknownTable>(); }},
		{ID_DATABASE_INCORRECT_PASSWORD, []() { return std::make_unique<DatabaseIncorrectPassword>(); }},
		{ID_READY_EVENT_SET, []() { return std::make_unique<ReadyEventSet>(); }},
		{ID_READY_EVENT_UNSET, []() { return std::make_unique<ReadyEventUnset>(); }},
		{ID_READY_EVENT_ALL_SET, []() { return std::make_unique<ReadyEventAllSet>(); }},
		{ID_READY_EVENT_QUERY, []() { return std::make_unique<ReadyEventQuery>(); }},
		{ID_LOBBY_GENERAL, []() { return std::make_unique<LobbyGeneral>(); }},
		{ID_AUTO_RPC_CALL, []() { return std::make_unique<AutoRPCCall>(); }},
		{ID_AUTO_RPC_REMOTE_INDEX, []() { return std::make_unique<AutoRPCRemoteIndex>(); }},
		{ID_AUTO_RPC_UNKNOWN_REMOTE_INDEX, []() { return std::make_unique<AutoRPCUnknownRemoteIndex>(); }},
		{ID_RPC_REMOTE_ERROR, []() { return std::make_unique<RPCRemoteError>(); }},
		{ID_USER_PACKET_ENUM, []() { return std::make_unique<UserPacketEnum>(); }},
	};

	void HandlePacket(RakNet::BitStream& inStream, const SystemAddress& sysAddr) {
		RaknetPacket data;
		if (!data.Deserialize(inStream)) {
			fprintf(stderr, "Error reading packet header\n");
			return;
		}

		auto it = g_Handlers.find(data.messageID);
		if (it != g_Handlers.end()) {
			auto packet = it->second();
			packet->sysAddr = sysAddr;
			if (!packet->Deserialize(inStream)) {
				fprintf(stderr, "Error reading packet data for message ID: %i\n", data.messageID);
				return;
			}
			packet->Handle();
		} else {
			fprintf(stderr, "Unhandled packet with ID: %i\n", data.messageID);
		}
	}
}