//###########################################################################
// btl4console -- the in-process LocalConsole marshal (BT_GLASS only;
// compiled only when the gate is on -- see CMakeLists.txt).
// Protocol + design: btl4console.hpp and tools/btconsole.py (the wire
// ground truth; every constant below matches it).
//###########################################################################

#include "btl4console.hpp"

#include <winsock2.h>
#include <ws2tcpip.h>
#include <windows.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <stdarg.h>

//
// The marshal logs to its own file (marshal.log in the cwd = content\):
// the engine's DEBUG log stream is single-threaded and this is a worker
// thread.  Ordinary printf-file logging, flushed per line.
//
static FILE *marshalLogFile = NULL;
static void
	MarshalLog(const char *format, ...)
{
	if (marshalLogFile == NULL)
	{
		marshalLogFile = fopen("marshal.log", "at");
		if (marshalLogFile == NULL)
		{
			return;
		}
	}
	va_list arguments;
	va_start(arguments, format);
	fprintf(marshalLogFile, "[marshal] ");
	vfprintf(marshalLogFile, format, arguments);
	fprintf(marshalLogFile, "\n");
	fflush(marshalLogFile);
	va_end(arguments);
}

//
// Wire constants -- tools/btconsole.py (T2-verified via BT_NET_PROBE=1).
//
enum
{
	ChunkBytes             = 1000,
	MessageLength          = 1024,	// sizeof(ReceiveEggFileMessage)
	EggMessageID           = 3,		// NetworkManager::ReceiveEggFileMessageID
	ReliableFlag           = 1,
	ConsoleHostID          = 1,		// FirstLegalHostID
	ApplicationClientID    = 4,		// NetworkClient::ApplicationClientID
	RunMissionMessageID    = 5,		// Application::RunMissionMessageID
	StopMissionMessageID   = 6,		// Application::StopMissionMessageID (enum: Run+1, APP.h:383)
	LaunchSettleSeconds    = 20,	// egg -> first launch (WaitingForLaunch)
	LaunchStepSeconds      = 4,		// first -> second launch (Launching -> Running)
	MissionEndGraceSeconds = 8		// stop sent -> relaunch the menu
};

struct MarshalState
{
	char	eggPath[128];
	char	podList[512];
	int		missionSeconds;
};

static MarshalState marshalState;

//
// The Steam-wire branch (BT_STEAM): a pod whose console address is a
// FakeIP is reached through the engine's Steam transport (munga_engine)
// instead of raw Winsock -- same wire bytes, different carrier.
//
#ifdef BT_STEAM
extern int    BTSteamNet_IsFakeAddress(unsigned long internet_address_be);
extern SOCKET BTSteamNet_Connect(unsigned long internet_address_be, int remote_port);
extern int    BTSteamNet_Owns(SOCKET wire_socket);
extern int    BTSteamNet_Send(SOCKET wire_socket, const char *data, int length);
extern int    BTSteamNet_Recv(SOCKET wire_socket, char *buffer, int capacity);
extern void   BTSteamNet_Close(SOCKET wire_socket);
#endif

static int
	MarshalIsSteam(SOCKET s)
{
#ifdef BT_STEAM
	return BTSteamNet_Owns(s);
#else
	(void)s;
	return 0;
#endif
}

static int
	MarshalSend(SOCKET s, const char *data, int length)
{
#ifdef BT_STEAM
	if (BTSteamNet_Owns(s))
	{
		return BTSteamNet_Send(s, data, length);
	}
#endif
	return send(s, data, length, 0);
}

static void
	MarshalClose(SOCKET s)
{
#ifdef BT_STEAM
	if (BTSteamNet_Owns(s))
	{
		BTSteamNet_Close(s);
		return;
	}
#endif
	closesocket(s);
}

//###########################################################################

void
	BTFE_RelaunchSelfAndExit(const char *arguments)
{
	//
	// A MENU relaunch (empty arguments) must not leak the mission
	// process's marshal handoff into the child -- the child inherits our
	// environment, and a stale BT_FE_EGG would skip the menu and arm a
	// marshal with no pod listening (found live: the first full cycle
	// re-armed instead of returning to the menu).
	//
	if (arguments == NULL || arguments[0] == 0)
	{
		SetEnvironmentVariableA("BT_FE_EGG", NULL);
		SetEnvironmentVariableA("BT_FE_PODS", NULL);
		SetEnvironmentVariableA("BT_FE_SECS", NULL);
		SetEnvironmentVariableA("BT_FE_LOOP", NULL);
	}

	WCHAR exe_path[MAX_PATH];
	GetModuleFileNameW(NULL, exe_path, MAX_PATH);

	WCHAR command_line[MAX_PATH + 256];
	int n = 0;
	command_line[n++] = L'"';
	for (WCHAR *s = exe_path; *s; ++s) command_line[n++] = *s;
	command_line[n++] = L'"';
	if (arguments != NULL && arguments[0])
	{
		command_line[n++] = L' ';
		for (const char *s = arguments; *s && n < MAX_PATH + 250; ++s)
		{
			command_line[n++] = (WCHAR)*s;
		}
	}
	command_line[n] = 0;

	STARTUPINFOW startup;
	PROCESS_INFORMATION process;
	memset(&startup, 0, sizeof(startup));
	startup.cb = sizeof(startup);
	memset(&process, 0, sizeof(process));
	if (CreateProcessW(exe_path, command_line, NULL, NULL, FALSE,
		0, NULL, NULL, &startup, &process))
	{
		CloseHandle(process.hThread);
		CloseHandle(process.hProcess);
	}
	ExitProcess(0);
}

//###########################################################################
// Wire builders
//###########################################################################

//
// Egg text -> the wire image NotationFile::ReadText expects: NUL-separated
// lines (NOTATION.cpp walks strchr(buffer,0)+1 per line).
//
static unsigned char *
	EggWireImage(const char *path, int *out_length)
{
	FILE *f = fopen(path, "rb");
	if (f == NULL)
	{
		return NULL;
	}
	fseek(f, 0, SEEK_END);
	long raw_length = ftell(f);
	fseek(f, 0, SEEK_SET);
	char *raw = (char *)malloc(raw_length + 1);
	fread(raw, 1, raw_length, f);
	fclose(f);
	raw[raw_length] = 0;

	unsigned char *wire = (unsigned char *)malloc(raw_length + 2);
	int w = 0;
	for (long i = 0; i < raw_length; ++i)
	{
		char c = raw[i];
		if (c == '\r')
		{
			continue;
		}
		wire[w++] = (c == '\n') ? 0 : (unsigned char)c;
	}
	if (w == 0 || wire[w - 1] != 0)
	{
		wire[w++] = 0;					// terminate the final line
	}
	free(raw);
	*out_length = w;
	return wire;
}

static void
	PutInt32(unsigned char *at, int value)
{
	memcpy(at, &value, 4);				// x86: little-endian, matching the wire
}

//
// One 1040-byte egg chunk: NetworkPacketHeader (clientID=0, gameID=0,
// fromHost=1, timeStamp=0) + ReceiveEggFileMessage (len=1024, id=3,
// flags=Reliable, seq, fileLen, thisLen, data[1000]).
//
static int
	SendEggChunks(SOCKET s, const unsigned char *wire, int wire_length)
{
	unsigned char packet[16 + MessageLength];
	int sequence = 0;
	for (int off = 0; off < wire_length; off += ChunkBytes, ++sequence)
	{
		int this_length = wire_length - off;
		if (this_length > ChunkBytes)
		{
			this_length = ChunkBytes;
		}
		memset(packet, 0, sizeof(packet));
		PutInt32(packet + 0,  0);				// clientID (NetworkManager)
		PutInt32(packet + 4,  0);				// gameID
		PutInt32(packet + 8,  ConsoleHostID);	// fromHost
		PutInt32(packet + 12, 0);				// timeStamp
		PutInt32(packet + 16, MessageLength);
		PutInt32(packet + 20, EggMessageID);
		PutInt32(packet + 24, ReliableFlag);
		PutInt32(packet + 28, sequence);
		PutInt32(packet + 32, wire_length);
		PutInt32(packet + 36, this_length);
		memcpy(packet + 40, wire + off, this_length);
		if (MarshalSend(s, (const char *)packet, sizeof(packet)) != sizeof(packet))
		{
			return -1;
		}
	}
	return sequence;
}

//
// The 28-byte application message (RunMission / StopMission): header with
// clientID=ApplicationClientID + a 12-byte message (len=12, id, flags).
//
static int
	SendApplicationMessage(SOCKET s, int message_id)
{
	unsigned char packet[16 + 12];
	PutInt32(packet + 0,  ApplicationClientID);
	PutInt32(packet + 4,  0);
	PutInt32(packet + 8,  ConsoleHostID);
	PutInt32(packet + 12, 0);
	PutInt32(packet + 16, 12);
	PutInt32(packet + 20, message_id);
	PutInt32(packet + 24, ReliableFlag);
	return (MarshalSend(s, (const char *)packet, sizeof(packet)) == sizeof(packet))
		? 0 : -1;
}

//###########################################################################
// The marshal thread
//###########################################################################

static SOCKET
	ConnectWithRetry(const char *host, int port, int attempts)
{
#ifdef BT_STEAM
	{
		unsigned long address_be = inet_addr(host);
		if (address_be != INADDR_NONE && BTSteamNet_IsFakeAddress(address_be))
		{
			for (int attempt = 0; attempt < attempts; ++attempt)
			{
				SOCKET s = BTSteamNet_Connect(address_be, port);
				if (s != INVALID_SOCKET)
				{
					return s;
				}
				Sleep(1000);
			}
			return INVALID_SOCKET;
		}
	}
#endif
	for (int attempt = 0; attempt < attempts; ++attempt)
	{
		SOCKET s = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
		if (s == INVALID_SOCKET)
		{
			return INVALID_SOCKET;
		}
		sockaddr_in address;
		memset(&address, 0, sizeof(address));
		address.sin_family = AF_INET;
		address.sin_port = htons((unsigned short)port);
		address.sin_addr.s_addr = inet_addr(host);
		if (connect(s, (sockaddr *)&address, sizeof(address)) == 0)
		{
			int nodelay = 1;
			setsockopt(s, IPPROTO_TCP, TCP_NODELAY,
				(const char *)&nodelay, sizeof(nodelay));
			return s;
		}
		closesocket(s);
		Sleep(1000);					// the pod may not be listening yet
	}
	return INVALID_SOCKET;
}

static DWORD WINAPI
	MarshalThread(LPVOID)
{
	WSADATA wsa;
	WSAStartup(MAKEWORD(2, 2), &wsa);	// ref-counted; the engine does its own

	int wire_length = 0;
	unsigned char *wire = EggWireImage(marshalState.eggPath, &wire_length);
	if (wire == NULL)
	{
		MarshalLog("cannot read egg %s -- marshal down", marshalState.eggPath);
		return 1;
	}

	//
	// Connect every pod (self first -- our own engine is still booting;
	// retry like the real console).
	//
	SOCKET pods[8];
	char pod_names[8][64];
	int pod_count = 0;
	{
		char list_copy[512];
		strcpy(list_copy, marshalState.podList);
		char *cursor = strtok(list_copy, ",");
		while (cursor != NULL && pod_count < 8)
		{
			char host[48];
			int port = 0;
			const char *colon = strrchr(cursor, ':');
			if (colon != NULL && (port = atoi(colon + 1)) > 0 &&
				(int)(colon - cursor) < (int)sizeof(host))
			{
				memcpy(host, cursor, colon - cursor);
				host[colon - cursor] = 0;
				MarshalLog("connecting pod %s:%d", host, port);
				SOCKET s = ConnectWithRetry(host, port, 60);
				if (s == INVALID_SOCKET)
				{
					MarshalLog("pod %s:%d never answered -- aborting mission start",
						host, port);
					BTFE_RelaunchSelfAndExit("");
				}
				strcpy(pod_names[pod_count], cursor);
				pods[pod_count++] = s;
			}
			cursor = strtok(NULL, ",");
		}
	}
	if (pod_count == 0)
	{
		MarshalLog("empty pod list -- marshal down");
		return 1;
	}

	//
	// Stream the egg to every pod, settle, launch twice, hold the clock.
	//
	for (int p = 0; p < pod_count; ++p)
	{
		int chunks = SendEggChunks(pods[p], wire, wire_length);
		MarshalLog("egg -> %s (%d bytes, %d chunks)", pod_names[p],
			wire_length, chunks);
	}
	free(wire);

	Sleep(LaunchSettleSeconds * 1000);
	for (int p = 0; p < pod_count; ++p)
	{
		SendApplicationMessage(pods[p], RunMissionMessageID);
	}
	MarshalLog("RunMission #1 sent to %d pod(s)", pod_count);
	Sleep(LaunchStepSeconds * 1000);
	for (int p = 0; p < pod_count; ++p)
	{
		SendApplicationMessage(pods[p], RunMissionMessageID);
	}
	MarshalLog("RunMission #2 sent -- mission running; clock %ds",
		marshalState.missionSeconds);

	//
	// Own the mission clock.  Drain (and discard) pod->console traffic so
	// the sockets stay healthy -- the console must STAY CONNECTED (a
	// disconnect trips the pod's console-loss path).
	//
	DWORD mission_end = GetTickCount()
		+ (DWORD)marshalState.missionSeconds * 1000;
	char drain[4096];
	while (GetTickCount() < mission_end)
	{
		//
		// select() only understands REAL sockets; Steam pseudo-sockets
		// drain nonblocking on their own pass.
		//
		fd_set readable;
		FD_ZERO(&readable);
		int real_count = 0;
		for (int p = 0; p < pod_count; ++p)
		{
			if (!MarshalIsSteam(pods[p]))
			{
				FD_SET(pods[p], &readable);
				++real_count;
			}
		}
		if (real_count > 0)
		{
			timeval tv = { 1, 0 };
			int ready = select(0, &readable, NULL, NULL, &tv);
			if (ready > 0)
			{
				for (int p = 0; p < pod_count; ++p)
				{
					if (!MarshalIsSteam(pods[p]) &&
						FD_ISSET(pods[p], &readable))
					{
						recv(pods[p], drain, sizeof(drain), 0);
					}
				}
			}
		}
		else
		{
			Sleep(1000);
		}
#ifdef BT_STEAM
		for (int p = 0; p < pod_count; ++p)
		{
			if (MarshalIsSteam(pods[p]))
			{
				while (BTSteamNet_Recv(pods[p], drain, sizeof(drain)) > 0)
				{
				}
			}
		}
#endif
	}

	MarshalLog("mission clock expired -- StopMission to all pods");
	for (int p = 0; p < pod_count; ++p)
	{
		SendApplicationMessage(pods[p], StopMissionMessageID);
	}
	Sleep(MissionEndGraceSeconds * 1000);
	for (int p = 0; p < pod_count; ++p)
	{
		MarshalClose(pods[p]);
	}

	MarshalLog("relaunching the menu");
	BTFE_RelaunchSelfAndExit("");
	return 0;
}

int
	BTLocalConsole_Start(
		const char *egg_path,
		const char *pod_list,
		int mission_seconds)
{
	strncpy(marshalState.eggPath, egg_path, sizeof(marshalState.eggPath) - 1);
	strncpy(marshalState.podList, pod_list, sizeof(marshalState.podList) - 1);
	marshalState.missionSeconds =
		(mission_seconds > 0) ? mission_seconds : 600;

	HANDLE thread = CreateThread(NULL, 0, MarshalThread, NULL, 0, NULL);
	if (thread == NULL)
	{
		return -1;
	}
	CloseHandle(thread);
	return 0;
}