#include "mungal4.h"
#pragma hdrstop

#include "l4steamtransport.h"

#ifdef RP412_STEAM

#include "l4nettransport.h"

// The engine builds with /Zp1; Valve's ABI expects default packing.
// Callback structs carry their own pack pragmas, but everything else
// must see the platform default.
#pragma pack(push, 8)
#include "steam/steam_api.h"
#include "steam/isteamnetworkingsockets.h"
#include "steam/isteamnetworkingutils.h"
#include "steam/steamnetworkingfakeip.h"
#pragma pack(pop)

//########################################################################
// SteamNetTransport - see l4steamtransport.h for the design. The
// engine is single threaded and so is this: every entry point pumps
// SteamAPI_RunCallbacks, which is where the status-changed callback
// fires (accepting incoming connections and marking drops).
//########################################################################

namespace
{
	enum
	{
		maxListeners	= 4,
		maxConnections	= 32,
		maxPeers		= 16,
		maxPending		= 8,
		leftoverSize	= 4096
	};

	struct ListenerRecord
	{
		HSteamListenSocket	handle;
		int					fakePortIndex;
		unsigned short		enginePort;			// host order
		HSteamNetConnection	pending[maxPending];
		int					pendingCount;
	};

	struct ConnectionRecord
	{
		HSteamNetConnection	handle;
		unsigned long		remoteIP;			// network order
		unsigned short		remoteEnginePort;	// network order
		Logical				connected;
		Logical				dead;
		char				leftover[leftoverSize];
		int					leftoverCount;
	};

	struct PeerRecord
	{
		unsigned long		fakeIP;				// host order
		unsigned short		fakeGamePort;		// host order
	};

	Logical				gSteamReady = False;
	unsigned long		gLocalFakeIP = 0;		// host order
	unsigned short		gLocalFakePorts[2] = { 0, 0 };
	char				gLocalFakeAddressString[32] = "";

	// engine ports in the order Listen sees them: [0] console, [1] game
	unsigned short		gEnginePorts[2] = { 0, 0 };
	int					gEnginePortCount = 0;

	ListenerRecord		gListeners[maxListeners];
	int					gListenerCount = 0;
	ConnectionRecord	gConnections[maxConnections];
	int					gConnectionCount = 0;
	PeerRecord			gPeers[maxPeers];
	int					gPeerCount = 0;

	ListenerRecord *
		FindListener(HSteamListenSocket handle)
	{
		for (int i = 0; i < gListenerCount; ++i)
		{
			if (gListeners[i].handle == handle)
			{
				return &gListeners[i];
			}
		}
		return NULL;
	}

	ConnectionRecord *
		FindConnection(HSteamNetConnection handle)
	{
		for (int i = 0; i < gConnectionCount; ++i)
		{
			if (gConnections[i].handle == handle)
			{
				return &gConnections[i];
			}
		}
		return NULL;
	}

	ConnectionRecord *
		AddConnection(
			HSteamNetConnection handle,
			unsigned long remote_ip_net,
			unsigned short remote_engine_port_net
		)
	{
		if (gConnectionCount >= maxConnections)
		{
			return NULL;
		}
		ConnectionRecord *record = &gConnections[gConnectionCount++];
		memset(record, 0, sizeof(*record));
		record->handle = handle;
		record->remoteIP = remote_ip_net;
		record->remoteEnginePort = remote_engine_port_net;
		return record;
	}

	void
		RemoveConnection(HSteamNetConnection handle)
	{
		for (int i = 0; i < gConnectionCount; ++i)
		{
			if (gConnections[i].handle == handle)
			{
				gConnections[i] = gConnections[gConnectionCount - 1];
				--gConnectionCount;
				return;
			}
		}
	}

	unsigned short
		LookupPeerFakeGamePort(unsigned long fake_ip_host)
	{
		for (int i = 0; i < gPeerCount; ++i)
		{
			if (gPeers[i].fakeIP == fake_ip_host)
			{
				return gPeers[i].fakeGamePort;
			}
		}
		return 0;
	}

	//---------------------------------------------------------------
	// Status-changed callback (fires inside SteamAPI_RunCallbacks on
	// the game thread): accept incoming, queue connected, mark drops
	//---------------------------------------------------------------
	void __cdecl
		OnConnectionStatusChanged(SteamNetConnectionStatusChangedCallback_t *status)
	{
		switch (status->m_info.m_eState)
		{
		case k_ESteamNetworkingConnectionState_Connecting:
			if (status->m_info.m_hListenSocket != k_HSteamListenSocket_Invalid)
			{
				// incoming: accept immediately, queue it when Connected
				SteamNetworkingSockets()->AcceptConnection(status->m_hConn);
			}
			break;

		case k_ESteamNetworkingConnectionState_Connected:
			if (status->m_info.m_hListenSocket != k_HSteamListenSocket_Invalid)
			{
				ListenerRecord *listener = FindListener(status->m_info.m_hListenSocket);
				if (listener != NULL && listener->pendingCount < maxPending)
				{
					listener->pending[listener->pendingCount++] = status->m_hConn;
				}
			}
			else
			{
				ConnectionRecord *record = FindConnection(status->m_hConn);
				if (record != NULL)
				{
					record->connected = True;
				}
			}
			break;

		case k_ESteamNetworkingConnectionState_ProblemDetectedLocally:
		case k_ESteamNetworkingConnectionState_ClosedByPeer:
			{
				ConnectionRecord *record = FindConnection(status->m_hConn);
				if (record != NULL)
				{
					record->dead = True;
				}
			}
			break;
		}
	}

	//---------------------------------------------------------------
	// The transport
	//---------------------------------------------------------------
	class SteamNetTransport:
		public NetTransport
	{
	public:
		Logical
			Startup()
		{
			// Install() already brought Steam up; the network manager
			// just confirms the transport is live
			return gSteamReady;
		}

		void
			Cleanup()
		{
			// mission teardown (single-binary race loop): drop the
			// wire but keep the Steam API up - the lobby lives on
			for (int i = 0; i < gConnectionCount; ++i)
			{
				SteamNetworkingSockets()->CloseConnection(
					gConnections[i].handle, 0, "mission teardown", true);
			}
			gConnectionCount = 0;
			for (int j = 0; j < gListenerCount; ++j)
			{
				SteamNetworkingSockets()->CloseListenSocket(gListeners[j].handle);
			}
			gListenerCount = 0;
			gEnginePortCount = 0;
		}

		int
			GetLocalAddresses(
				unsigned long *addresses,
				int max_count
			)
		{
			if (!gSteamReady || max_count < 1)
			{
				return 0;
			}
			addresses[0] = htonl(gLocalFakeIP);
			return 1;
		}

		Logical
			Resolve(
				const char *host_name,
				SOCKADDR_IN *address
			)
		{
			memset(address, 0, sizeof(SOCKADDR_IN));
			address->sin_family = AF_INET;

			SteamNetworkingIPAddr parsed;
			if (!parsed.ParseString(host_name))
			{
				return False;
			}
			address->sin_addr.S_un.S_addr = htonl(parsed.GetIPv4());
			address->sin_port = htons(parsed.m_port);	// 0 when absent
			return True;
		}

		Connection
			Connect(
				const SOCKADDR_IN *remote,
				int local_port
			)
		{
			unsigned long remote_fake_ip = ntohl(remote->sin_addr.S_un.S_addr);
			unsigned short fake_game_port = LookupPeerFakeGamePort(remote_fake_ip);
			if (fake_game_port == 0)
			{
				DEBUG_STREAM << "SteamNetTransport: no registered peer for "
					<< inet_ntoa(remote->sin_addr) << " - lobby did not feed it\n" << std::flush;
				return InvalidConnection;
			}

			DEBUG_STREAM << "SteamNetTransport: connecting to "
				<< inet_ntoa(remote->sin_addr) << " (fake port "
				<< fake_game_port << ")...\n" << std::flush;

			SteamNetworkingConfigValue_t option;
			option.SetPtr(
				k_ESteamNetworkingConfig_Callback_ConnectionStatusChanged,
				(void *) &OnConnectionStatusChanged);

			SteamNetworkingIPAddr target;
			target.Clear();
			target.SetIPv4(remote_fake_ip, fake_game_port);

			// mirror the TCP retry-while-refused loop, bounded: the
			// egg-ACK ordering means the peer may not be listening yet
			DWORD deadline = GetTickCount() + 120 * 1000;
			for (;;)
			{
				HSteamNetConnection handle =
					SteamNetworkingSockets()->ConnectByIPAddress(target, 1, &option);
				if (handle == k_HSteamNetConnection_Invalid)
				{
					return InvalidConnection;
				}
				ConnectionRecord *record =
					AddConnection(handle, remote->sin_addr.S_un.S_addr, remote->sin_port);

				while (!record->connected && !record->dead &&
					(LONG)(GetTickCount() - deadline) < 0)
				{
					SteamAPI_RunCallbacks();
					Sleep(25);
				}
				if (record->connected)
				{
					return (Connection) handle;
				}

				// attempt failed - drop it and retry until the deadline
				SteamNetworkingSockets()->CloseConnection(handle, 0, "retry", false);
				RemoveConnection(handle);
				if ((LONG)(GetTickCount() - deadline) >= 0)
				{
					DEBUG_STREAM << "SteamNetTransport: connect timed out\n" << std::flush;
					return InvalidConnection;
				}
				Sleep(250);
			}
		}

		Connection
			Listen(
				int local_port,
				int backlog
			)
		{
			if (gListenerCount >= maxListeners)
			{
				return InvalidConnection;
			}

			// engine port -> fake port index, in order of appearance:
			// the console channel always listens first, the mesh second
			int index = -1;
			for (int i = 0; i < gEnginePortCount; ++i)
			{
				if (gEnginePorts[i] == (unsigned short) local_port)
				{
					index = i;
				}
			}
			if (index < 0)
			{
				if (gEnginePortCount >= 2)
				{
					DEBUG_STREAM << "SteamNetTransport: only two fake ports allocated!\n" << std::flush;
					return InvalidConnection;
				}
				index = gEnginePortCount;
				gEnginePorts[gEnginePortCount++] = (unsigned short) local_port;
			}

			DEBUG_STREAM << "SteamNetTransport: listening on engine port "
				<< local_port << " (fake port " << gLocalFakePorts[index]
				<< ")...\n" << std::flush;

			SteamNetworkingConfigValue_t option;
			option.SetPtr(
				k_ESteamNetworkingConfig_Callback_ConnectionStatusChanged,
				(void *) &OnConnectionStatusChanged);

			HSteamListenSocket handle =
				SteamNetworkingSockets()->CreateListenSocketP2PFakeIP(index, 1, &option);
			if (handle == k_HSteamListenSocket_Invalid)
			{
				return InvalidConnection;
			}

			ListenerRecord *listener = &gListeners[gListenerCount++];
			memset(listener, 0, sizeof(*listener));
			listener->handle = handle;
			listener->fakePortIndex = index;
			listener->enginePort = (unsigned short) local_port;
			return (Connection) handle;
		}

		Connection
			Accept(Connection listener_handle)
		{
			SteamAPI_RunCallbacks();

			ListenerRecord *listener = FindListener((HSteamListenSocket) listener_handle);
			if (listener == NULL || listener->pendingCount == 0)
			{
				return InvalidConnection;
			}

			HSteamNetConnection handle = listener->pending[0];
			for (int i = 1; i < listener->pendingCount; ++i)
			{
				listener->pending[i - 1] = listener->pending[i];
			}
			--listener->pendingCount;

			// the remote engine port is the same engine port we accept
			// on: all pods share the -net convention under Steam
			unsigned long remote_ip_net = 0;
			SteamNetworkingIPAddr remote_fake;
			if (SteamNetworkingSockets()->GetRemoteFakeIPForConnection(
					handle, &remote_fake) == k_EResultOK)
			{
				remote_ip_net = htonl(remote_fake.GetIPv4());
			}
			ConnectionRecord *record = AddConnection(
				handle, remote_ip_net, htons(listener->enginePort));
			if (record != NULL)
			{
				record->connected = True;
			}
			return (Connection) handle;
		}

		void
			Close(Connection connection)
		{
			ListenerRecord *listener = FindListener((HSteamListenSocket) connection);
			if (listener != NULL)
			{
				SteamNetworkingSockets()->CloseListenSocket(listener->handle);
				*listener = gListeners[gListenerCount - 1];
				--gListenerCount;
				return;
			}
			SteamNetworkingSockets()->CloseConnection(
				(HSteamNetConnection) connection, 0, "closed", true);
			RemoveConnection((HSteamNetConnection) connection);
		}

		int
			Send(
				Connection connection,
				const void *data,
				int size
			)
		{
			EResult result = SteamNetworkingSockets()->SendMessageToConnection(
				(HSteamNetConnection) connection, data, (uint32) size,
				k_nSteamNetworkingSend_Reliable, NULL);
			return (result == k_EResultOK) ? size : -1;
		}

		int
			Receive(
				Connection connection,
				void *buffer,
				int size
			)
		{
			SteamAPI_RunCallbacks();

			ConnectionRecord *record = FindConnection((HSteamNetConnection) connection);
			if (record == NULL)
			{
				return ReceiveNoData;
			}

			char *out = (char *) buffer;
			int copied = 0;

			// leftover bytes from a message that outsized the last call
			if (record->leftoverCount > 0)
			{
				int take = (record->leftoverCount < size) ? record->leftoverCount : size;
				memcpy(out, record->leftover, take);
				memmove(record->leftover, record->leftover + take,
					record->leftoverCount - take);
				record->leftoverCount -= take;
				out += take;
				copied += take;
			}

			// drain messages while they fit; stash any partial tail
			while (copied < size && record->leftoverCount == 0)
			{
				SteamNetworkingMessage_t *message = NULL;
				int count = SteamNetworkingSockets()->ReceiveMessagesOnConnection(
					record->handle, &message, 1);
				if (count <= 0)
				{
					break;
				}
				int room = size - copied;
				int take = ((int) message->m_cbSize < room) ? (int) message->m_cbSize : room;
				memcpy(out, message->m_pData, take);
				out += take;
				copied += take;

				int rest = (int) message->m_cbSize - take;
				if (rest > 0)
				{
					if (rest > leftoverSize)
					{
						rest = leftoverSize;	// cannot happen: engine packets < 2K
					}
					memcpy(record->leftover, (const char *) message->m_pData + take, rest);
					record->leftoverCount = rest;
				}
				message->Release();
			}

			if (copied > 0)
			{
				return copied;
			}
			if (record->dead)
			{
				return ReceiveDisconnected;
			}
			return ReceiveNoData;
		}

		Logical
			GetRemoteAddress(
				Connection connection,
				SOCKADDR_IN *address
			)
		{
			memset(address, 0, sizeof(SOCKADDR_IN));
			ConnectionRecord *record = FindConnection((HSteamNetConnection) connection);
			if (record == NULL || !record->connected)
			{
				return False;
			}
			address->sin_family = AF_INET;
			address->sin_addr.S_un.S_addr = record->remoteIP;
			address->sin_port = record->remoteEnginePort;
			return True;
		}
	};

	SteamNetTransport gSteamTransport;
}

//########################################################################
// Install: bring Steam up, get our FakeIP identity, take over the wire
//########################################################################

Logical
	SteamNetTransport_Install()
{
	if (gSteamReady)
	{
		return True;
	}

	if (!SteamAPI_Init())
	{
		DEBUG_STREAM << "SteamNetTransport: SteamAPI_Init failed "
			<< "(Steam not running, or no steam_appid.txt) - staying on TCP\n" << std::flush;
		return False;
	}

	SteamNetworkingUtils()->InitRelayNetworkAccess();

	if (!SteamNetworkingSockets()->BeginAsyncRequestFakeIP(2))
	{
		DEBUG_STREAM << "SteamNetTransport: BeginAsyncRequestFakeIP failed - staying on TCP\n" << std::flush;
		return False;
	}

	// FakeIP allocation is async: pump until it lands (or give up)
	SteamNetworkingFakeIPResult_t fake;
	memset(&fake, 0, sizeof(fake));
	DWORD deadline = GetTickCount() + 20 * 1000;
	for (;;)
	{
		SteamAPI_RunCallbacks();
		SteamNetworkingSockets()->GetFakeIP(0, &fake);
		if (fake.m_eResult == k_EResultOK)
		{
			break;
		}
		if (fake.m_eResult != k_EResultBusy && fake.m_eResult != k_EResultNoMatch)
		{
			DEBUG_STREAM << "SteamNetTransport: FakeIP allocation failed (EResult "
				<< (int) fake.m_eResult << ") - staying on TCP\n" << std::flush;
			return False;
		}
		if ((LONG)(GetTickCount() - deadline) >= 0)
		{
			DEBUG_STREAM << "SteamNetTransport: FakeIP allocation timed out - staying on TCP\n" << std::flush;
			return False;
		}
		Sleep(50);
	}

	gLocalFakeIP = fake.m_unIP;
	gLocalFakePorts[0] = fake.m_unPorts[0];
	gLocalFakePorts[1] = fake.m_unPorts[1];

	SteamNetworkingIPAddr mine;
	mine.Clear();
	mine.SetIPv4(gLocalFakeIP, 0);
	mine.ToString(gLocalFakeAddressString, sizeof(gLocalFakeAddressString), false);

	DEBUG_STREAM << "SteamNetTransport: up as " << gLocalFakeAddressString
		<< " (fake ports " << gLocalFakePorts[0] << " console, "
		<< gLocalFakePorts[1] << " game)\n" << std::flush;

	gSteamReady = True;
	NetTransport_Set(&gSteamTransport);
	return True;
}

const char *
	SteamNetTransport_GetFakeAddressString()
{
	return gLocalFakeAddressString;
}

int
	SteamNetTransport_GetFakeGamePort()
{
	return gLocalFakePorts[1];
}

void
	SteamNetTransport_RegisterPeer(
		const char *fake_ip,
		int fake_game_port
	)
{
	if (gPeerCount >= maxPeers)
	{
		return;
	}
	SteamNetworkingIPAddr parsed;
	if (!parsed.ParseString(fake_ip))
	{
		return;
	}
	gPeers[gPeerCount].fakeIP = parsed.GetIPv4();
	gPeers[gPeerCount].fakeGamePort = (unsigned short) fake_game_port;
	++gPeerCount;
	DEBUG_STREAM << "SteamNetTransport: peer registered: " << fake_ip
		<< " game port " << fake_game_port << "\n" << std::flush;
}

#endif	// RP412_STEAM