/* Copyright (c) 2006-2007, H. Andres Lagar-Cavilla, University of Toronto
 * In the context of the vmgl project
 * Module in charge of resuming suspended OpenGL state
 *
 * See the file LICENSE.txt for information on redistributing this software.
 */
 
/* High-level overview of what's going on here.
 * We use the state tracker diff functions to send OpenGL state to
 * a new stub. We diff against an empty context and we ship the diffs.
 * We use the Display List Manager (DLM) code to keep unrolled display lists
 * in memory. We then send them to the stub if need be.
 * The new autogenerated set of special functions (packspu_resume_autogen.py)
 * do the appropriate packing, DLM accounting, and if need be leave traces of 
 * their existence in the state tracker code. Existing packspu special 
 * functions have been modified likewise. We create a Diff Dispatch Table
 * (diffgen.py) that simply shortcuts calls into the packer. This Diff DT is 
 * passed to context diffing, texture diffing, and DLM replay functions so that
 * OpenGL state gets actually sent to the new stub.
 * Finally, when replaying a display list, the DLM will also apply all changes 
 * the state tracker. The DLMReplayCallList(s) functions get passed the State 
 * Dispatch Table (stategen.py) which simply shortcuts to correpsonding 
 * crState calls.
 */

#include "packspu.h"
#include "cr_environment.h"
#include "cr_error.h"
#include "cr_mem.h"
#include "cr_string.h"
#include "cr_packfunctions.h"
#include "cr_url.h"
#include "cr_process.h"
#include "cr_XExtension.h"

#include <string.h>
#include <errno.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/un.h>
#include <unistd.h>
#include <sys/time.h>

static CRHashTable *oldHash;

/* Return one on success, zero on failure. Establish new TCP connection,
 * replace data of old connection, reset packer settings. */
static int connectToNewServer(char *glStub_str) {
	GET_THREAD(thread);
	CRNetServer netServer;
	int i;
	CRConnection *newConn = NULL;

	if (!thread) {
	    crInfo("Resume: can't resume if no pack spu threads allocated!");
	    free(glStub_str);
	    return 0;
	}
	netServer = thread->netServer;

	crDebug("Resume: connecting to new glStub daemon %s", glStub_str);
	
	setUpServerConnection(glStub_str); /* glStub_str is freed inside this */

	crDebug("Resume: connecting to new glStub %s", pack_spu.name);

	/* Reset netServer struct vals */
	crFree(netServer.name);
	netServer.name = crStrdup( pack_spu.name );
	netServer.buffer_size = pack_spu.buffer_size;

	/* Attempt to reconnect to the glStub instance. Store 
	 * connection in temp var */
        for (i=0; i<PACKSPU_CONNECT_RETRIES; i++) {
    	    newConn = crNetConnectToServer( netServer.name, DEFAULT_SERVER_PORT, netServer.buffer_size );
            if (newConn) break;
            crMsleep(PACKSPU_CONNECT_RETRY_SLEEP_MS);
        }
	if (!newConn) {
		crWarning("pack SPU: failed to connect to %s", pack_spu.name);
		return 0;
	}

	/* Copy new conn into old conn. Uber hack. No mem leaks on
	 * netServer.conn, apparently (hostname got freed earlier). 
	 * On return from attemptResume, everyone will be fooled. */
	memcpy(netServer.conn, newConn, sizeof(CRConnection));

	/* Packer and miscellaneous setup based on new conn. Hack */
	pack_spu.swap = netServer.conn->swap;
        thread->packer = crPackNewContext( pack_spu.swap );
        CRASSERT(thread->packer);
        thread->buffer.context = NULL;
        crPackInitBuffer( &(thread->buffer), thread->buffer.pack,
                            netServer.conn->buffer_size, netServer.conn->mtu );
        thread->buffer.canBarf = netServer.conn->Barf ? GL_TRUE : GL_FALSE;
        crPackSetBuffer( thread->packer, &thread->buffer );
        crPackFlushFunc( thread->packer, packspuFlush );
        crPackFlushArg( thread->packer, (void *) thread );
        crPackSendHugeFunc( thread->packer, packspuHuge );
        crPackSetContext( thread->packer );

	return 1;
}

/**** Window resume code here ****/

/* We get here by walking the old XWindows hash table.
 * Instruct glStub to create a window like the one we had (same visual).
 * Re-enter it into the new hash table with the new glStub-provided id.
 * Re-register it with the VMGLExtension on the X server. */
static void packspuResumeWindow(unsigned long key, void *data1, void *data2)
{
	char fakeDpy[25];
	WindowInfo *newWindow;

	ThreadInfo *thread = (ThreadInfo *) data2;
	WindowInfo *winInfo = (WindowInfo *) data1;
	GLint window = 0;
	GLint writeback = 1;
	
	crDebug("Resume: handling X window %d", winInfo->XWindow);

	/* Instruct glStub to create a window like the one we had. */
	if (pack_spu.swap)
		crPackWindowCreateSWAP( fakeDpy, winInfo->visual, &window, &writeback);
	else
		crPackWindowCreate( fakeDpy, winInfo->visual, &window, &writeback);
	packspuFlush(thread);
	while (writeback) {
		crNetRecv();
	}
	if (pack_spu.swap)
		window = (GLint) SWAP32(window);
	if (window < 0) {
		crWarning("Resume: can't recreate XWindow %d.", winInfo->XWindow);
		return;
	}
	
	/* Re-enter it into the new hash table with the new glStub side id */
        newWindow = (WindowInfo *) crAlloc(sizeof(WindowInfo));
        newWindow->XWindow = winInfo->XWindow;
        newWindow->visual = winInfo->visual;
        crHashtableAdd(pack_spu.XWindows, window, newWindow);
                                                                	
	/* Re-register it with the VMGLExtension on the X server.
	 * The extension will take care of window position, size and visibility. */
	if (pack_spu.openedXDisplay == XDPY_CONNECTED)
    	    if (!(XVMGLExtWatchWindow(pack_spu.XDisplay, window, newWindow->XWindow)))
		crInfo("XVMGLExtWatchWindow for window %d failed (re-registering %d X Window).", 
			window, newWindow->XWindow);
	
	/* Uber hack to address windows by their old id and obtaining the new one */
	winInfo->XWindow = window;

}

/**** Context resume: the tough meat ****/
 
/* Call back used to walk the texture hash table for a given context */
static void TextureCallback( unsigned long key, void *data1, void *data2 )
{
	CRbitvalue *b = NULL, *n = NULL;
	CRContext *context = (CRContext *) data2;
	CRTextureObj *texture = (CRTextureObj *) data1;

	if (!texture) return;
	crStateTextureObjectDiff(context, b, n, texture, GL_TRUE);
}


/* This does one third of context resume, the textures.
 * We take advantage of the hash table of existing textures. 
 * We walk it and use crStateTextureObjectDiff with dumb parameters
 * to obtain the whole texture. The diff_DT functions pack and send 
 * that data. */
static void packspuResumeTextures(CRContext *tmpContext, CRContext *state)
{
	CRTextureState *textureState = &(state->texture);
	CRbitvalue *b = NULL, *n = NULL;

	crDebug("Resume: Textures");

	/* Set active texture to unit 0 and walk the texture table*/
	if (pack_spu.swap)
		crPackActiveTextureARBSWAP(GL_TEXTURE0);
	else
		crPackActiveTextureARB(GL_TEXTURE0);
	crHashtableWalk(state->shared->textureTable, TextureCallback, tmpContext);

	/* Special-cased base textures, essentially cut&paste from the above callback. */
	crStateTextureObjectDiff(tmpContext, b, n, &textureState->baseCubeMap, GL_TRUE);
	crStateTextureObjectDiff(tmpContext, b, n, &textureState->baseRect, GL_TRUE);
	crStateTextureObjectDiff(tmpContext, b, n, &textureState->base1D, GL_TRUE);
	crStateTextureObjectDiff(tmpContext, b, n, &textureState->base2D, GL_TRUE);
	crStateTextureObjectDiff(tmpContext, b, n, &textureState->base3D, GL_TRUE);

	/* After walking thorugh all tables, need to reinstate data for active texture in unit 0 */
	if (pack_spu.swap) {
		crPackActiveTextureARBSWAP(GL_TEXTURE0);
		crPackBindTextureSWAP(GL_TEXTURE_CUBE_MAP_ARB, textureState->unit[0].currentTextureCubeMap->name);
		crPackBindTextureSWAP(GL_TEXTURE_RECTANGLE_NV, textureState->unit[0].currentTextureRect->name);
		crPackBindTextureSWAP(GL_TEXTURE_1D, textureState->unit[0].currentTexture1D->name);
		crPackBindTextureSWAP(GL_TEXTURE_2D, textureState->unit[0].currentTexture2D->name);
		crPackBindTextureSWAP(GL_TEXTURE_3D, textureState->unit[0].currentTexture3D->name);
	}
	else {
		crPackActiveTextureARB(GL_TEXTURE0);
		crPackBindTexture(GL_TEXTURE_CUBE_MAP_ARB, textureState->unit[0].currentTextureCubeMap->name);
		crPackBindTexture(GL_TEXTURE_RECTANGLE_NV, textureState->unit[0].currentTextureRect->name);
		crPackBindTexture(GL_TEXTURE_1D, textureState->unit[0].currentTexture1D->name);
		crPackBindTexture(GL_TEXTURE_2D, textureState->unit[0].currentTexture2D->name);
		crPackBindTexture(GL_TEXTURE_3D, textureState->unit[0].currentTexture3D->name);
	}

	/* Done, set the active texture unit back */
	if (pack_spu.swap)
	    crPackActiveTextureARBSWAP(GL_TEXTURE0 + textureState->curTextureUnit);
	else
	    crPackActiveTextureARB(GL_TEXTURE0 + textureState->curTextureUnit);
}


/* This does another third of context resume, the display lists
 * which required pervasive DLM code insertions. Luckily, after
 * everything has been put in the DLM, we require a simple call
 * to DLMSendAllDLMLists */
static void packspuResumeLists(CRContext *tmpContext, CRDLM *DLM)
{
	crDebug("Resume: Display Lists");

	/* We setup pixel packing before sending the lists,
	 * and restore back to what it was after sending the lists. */
	crDLMSetupClientState(&pack_spu.diff_DT);
	crDLMSendAllDLMLists(DLM, &pack_spu.diff_DT);
	crDLMRestoreClientState(&tmpContext->client, &pack_spu.diff_DT);
}


/* This is a callback used by the contextList walking funtion
 * invoked by the main resume routine. It does the first third 
 * of context resume, the actual context data.
 * It then invokes the functions above to handle textures and DL's. */
static void packspuResumeContext(void *element, void *arg)
{
	ContextInfo *context;
	GLint shareCtx, return_val, serverWindow, nativeWindow;
	int writeback;
	CRContext *tmpContext;
	WindowInfo *winInfo;
	char fakeDpy[25];

	ThreadInfo *thread = (ThreadInfo *) arg;
	GLint sharedServerCtx = (GLint) -1;
	int slot = (int) element;

	/* Get this context */
        if ((slot < 0) || (slot >= pack_spu.numContexts) || (!(pack_spu.context[slot].clientState))) {
    	    crDebug("Resume: odd, context %d doesn't seem to be valid", slot + MAGIC_OFFSET);
    	    return;
        }
	context = &(pack_spu.context[slot]);
                                                                	
	crDebug("Resume: Context %d", slot + MAGIC_OFFSET);

	/* Find the shared context, if any */
	shareCtx = context->shareCtx;
	if (shareCtx > 0) {
	    shareCtx -= MAGIC_OFFSET;
	    if (shareCtx >= 0 && shareCtx < pack_spu.numContexts)
		sharedServerCtx = pack_spu.context[shareCtx].serverCtx;
	}
                                                                                                                                               
	/* Create the context on the new glStub */
	if (pack_spu.swap)
	    crPackCreateContextSWAP( fakeDpy, context->visual, sharedServerCtx, &return_val, &writeback);
	else
	    crPackCreateContext( fakeDpy, context->visual, sharedServerCtx, &return_val, &writeback);
	packspuFlush(thread);
	writeback = 1;
	while (writeback)
		crNetRecv();
	if (pack_spu.swap)
		return_val = (GLint) SWAP32(return_val);
	if (return_val <= 0) {
		crWarning("Resume: can't recreate context %d", slot + MAGIC_OFFSET);
		return;
	}
	context->serverCtx = return_val;


	/* Make the fresh context current on the glStub. */
	if (context->boundWindow) {
	    winInfo = crHashtableSearch(oldHash, context->boundWindow);
	    serverWindow = winInfo->XWindow;
	    winInfo = crHashtableSearch(pack_spu.XWindows, serverWindow);
	    nativeWindow = winInfo->XWindow;
	} else {
	    serverWindow = 0;
	    nativeWindow = 0;
	}
	if (pack_spu.swap)
	    crPackMakeCurrentSWAP(serverWindow, nativeWindow, return_val );
	else
	    crPackMakeCurrent(serverWindow, nativeWindow, return_val );

	/* This is the actual context state resume code. */
	/* Create a dummy context with pristine state. By diffing against 
	 * it we ensure we send everything. */
	tmpContext = crStateCreateContext(NULL, context->visual, NULL);
	/* Send context state */ 
	crStateDiffContext(tmpContext, context->clientState);
	/* Send texture state */
	packspuResumeTextures(tmpContext, context->clientState);
	/* Send display lists */
	packspuResumeLists(tmpContext, context->DLM);
	/* Destroy the temporary context, no longer needed */
	crStateDestroyContext(tmpContext);
	
	/* Done */
	packspuFlush(thread);
}
/**** Context resume code ends ****/

/* This is what this whole mess supports. GL state resumption.
 * Returns one on success (even partial), zero on failure to connect */
static int packspuResume(char *glStub_str) 
{
	GET_THREAD(thread);

	crDebug("Resume: beginning");
	
	/* First try to establish a connection with the server URL we've been given. */
	if (!(connectToNewServer(glStub_str))) /* glStub_str is freed and openedXDisplay set. */
	{
	    crDebug("Resume: Failed to reconnect!");
	    return 0;
	}
	
	/* Only attempt to restart if we are meant to use the X extension. */
	if (pack_spu.openedXDisplay != XDPY_DONT_CONNECT) 
	    openXDisplay("");

	/* Create windows first. Order irrelevant. Bit of a hack here.
	 * The WindowInfo structures are keyed into the hash table by glStub id.
	 * This id will naturally change now, so we need to re-key them into a new hash.
	 * We keep the old one in a temp var, and walk it to build the new one.
	 * The old hash table remains with the new glStub ids keyed by the old ones.
	 * We'll need that later. */
	crDebug("Resume: Windows first");
	oldHash = pack_spu.XWindows;
	pack_spu.XWindows = crAllocHashtable();
	crHashtableWalk(oldHash, packspuResumeWindow, thread);

	/* Now recreate contexts. Need to preserve ordering for shared context
	 * dependencies. Hence use list. */
	crDebug("Resume: Contexts second");
	crListApply(pack_spu.contextList, packspuResumeContext, thread);
	
	/* Restore current context. It was changed constantly as we walked the list. */
	if (thread->currentContext) {
		int resumeCurrentContext = (int) thread->currentContext->serverCtx;
		int resumeCurrentWindow_oldid = (int) thread->currentContext->boundWindow;
		WindowInfo *winInfo = crHashtableSearch(oldHash, resumeCurrentWindow_oldid);
		GLint resumeCurrentWindow = winInfo->XWindow;
		winInfo = crHashtableSearch(oldHash, resumeCurrentWindow); 
		GLint nativeWindow = winInfo->XWindow;
		
		crDebug("Resume: set back current context");
		crStateMakeCurrent(thread->currentContext->clientState );
                crDLMSetCurrentState(thread->currentContext->DLMState);
                thread->currentContext->boundWindow = resumeCurrentWindow;

		if (pack_spu.swap)
			crPackMakeCurrentSWAP(resumeCurrentWindow, nativeWindow, resumeCurrentContext);
	 	else
			crPackMakeCurrent(resumeCurrentWindow, nativeWindow, resumeCurrentContext);
	}

	/* Done */
	crDebug("Resume: ending");
	crFreeHashtable(oldHash, NULL);
	packspuFlush(thread);

	return 1;
}

/* Why isn this in the standard libraries? */
static inline unsigned long long timeval_diff(struct timeval start, struct timeval end) {
    unsigned long long micros = ((unsigned long long) end.tv_usec) - ((unsigned long long) start.tv_usec);
    unsigned long long seconds = ((unsigned long long) end.tv_sec) - ((unsigned long long) start.tv_sec);
    if (micros < 0) seconds ++;
    return (seconds*1000000) + micros;
}

/**** Entry point into this module ****/
void attemptResume() {
	int sock, len;
	struct sockaddr_un addr;
	char *glStub_str;
	size_t string_len, int_len = sizeof(size_t);
	ssize_t just_read;
	unsigned int pos;
	struct timeval susp_start, susp_end;
    
	/* First reset the state the VMGLExtension keeps for
	 * us on the X server. Only reset if the connection was open */
        if (pack_spu.openedXDisplay == XDPY_CONNECTED) {
    	    XVMGLExtReset(pack_spu.XDisplay);
	}
	pack_spu.openedXDisplay = XDPY_DONT_CONNECT; 
    	/* connectToNewServer in packsupResume will set it right */
    
        /* Now find out if we ought to resume and who is the new glStub. */
        crDebug("Resume: attempting to find new glStub to resume");

	/* Set up the unix domain socket */
	memset(&addr, 0, sizeof(struct sockaddr_un));
        addr.sun_family = AF_UNIX;
	sprintf(addr.sun_path, "%s/.vmgl_resume", crGetenv("HOME"));
	len = strlen(addr.sun_path) + sizeof(addr.sun_family);

	/* Loop endlessly, this could be a thin client configuration. */
	while(1) {
	    /* Socket calls */
	    sock = socket(AF_UNIX, SOCK_STREAM, 0);
	    if (sock == -1)
		crError("Resume: can't create UNIX domain socket: %s", strerror(errno));
	    if ((connect(sock, (struct sockaddr *) &addr, len)) == -1) {
		crDebug("Resume: connect to UNIX domain socket %s failed: %s",
			addr.sun_path, strerror(errno));
		close(sock);
		crMsleep(PACKSPU_CONNECT_RETRY_SLEEP_MS*2);
		continue;
	    }

	    /* Connected, get new glStub string */
	    if ((read(sock, (void *) &string_len, int_len)) < (ssize_t) int_len) {
		crDebug("Resume: couldn't even read an integer out of UNIX domain socket %s",
			addr.sun_path);
		close(sock);
		continue;
	    }
	    glStub_str = (char *) malloc(sizeof(char)*(string_len+1));
	    /* Kinda read exact loop */
	    pos = 0;
	    while (string_len > 0) {
		just_read = read(sock, (void *) &(glStub_str[pos]), string_len);
		if (just_read == -1) {
		    crDebug("Resume: error %s while reading from UNIX domain socket %s",
			    strerror(errno), addr.sun_path);
		    free(glStub_str);
		    glStub_str = NULL;
		    break;
		}
		string_len -= (size_t) just_read;
		pos += (unsigned int) just_read;
	    }
	    
	    /* Finally, invoke resume code */
	    close(sock);
	    if (glStub_str) {
		glStub_str[pos] = '\0';
		gettimeofday(&susp_start, NULL);
		/* glStub_str freed within call below */
		if (packspuResume(glStub_str)) {
		    gettimeofday(&susp_end, NULL);
		    crInfo("Resume: resume completed, connected to %s\n"
			   "Suspended state size %d, resume time in ms %llu\n",
			    pack_spu.name, pack_spu.thread[0].netServer.conn->total_bytes_sent,
			    (timeval_diff(susp_start, susp_end) / 1000));
		    break;
		} else
		    crDebug("Resume: failed attempt at resuming");
	    }    
	}
    
	/* We're out, means we're good */
}

