/* $Id: $ */

#include "config_nws.h"

#include <sys/types.h>
#include <signal.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <ctype.h>

#include "diagnostic.h"
#include "dnsutil.h"
#include "host_protocol.h"
#include "messages.h"
#include "osutil.h"
#include "strutil.h"
#include "nws_messages.h"
#include "register.h"
#include "nws_nameserver.h"
#include "timeouts.h"

typedef struct {
	char *name;
	Socket sd;
	int companion;
} reflectSocket;
#define HOW_MANY_IN_LIST	1024
static reflectSocket *myList;
static Socket *sockList;
static Socket proxyEar;


static int
FindSocket(	Socket sd,
		Socket *sList, 
		unsigned int howMany) {
	int i;

	/* sanity check */
	if (sd == NO_SOCKET || sList == NULL) {
		INFO("FindSocket: invalid parameter\n");
		return -1;
	}

	/* search for the socket */
	for (i = 0; i < howMany; i++) {
		if (sList[i] != NO_SOCKET) {
			break;
		}
	}
	if (i >= howMany) {
		i = -1;
	}

	return i;
}

static int
FindPairedSocket(	Socket sd,
			reflectSocket *r,
			unsigned int howMany) {
	int i;

	/* sanity check */
	if (sd == NO_SOCKET || r == NULL) {
		INFO("FindPairedSocket: invalid parameter\n");
		return -1;
	}

	/* find the socket */
	for (i = 0; i < howMany; i++) {
		if (r[i].sd == sd) {
			/* found */
			break;
		}
	}
	if (i >= howMany) {
		/* didn't find it */
		i = -1;
	}

	return i;
}


static int
PairSockets(	Socket one,
		Socket two,
		reflectSocket *r,
		unsigned int howMany) {
	int i, j;

	/* let's look for 2 free slots in the list */
	for (i = 0; i < howMany; i++) {
		if (r[i].sd == NO_SOCKET) {
			break;
		}
	}
	for (j = i + 1; j < howMany; j++) {
		if (r[j].sd == NO_SOCKET) {
			break;
		}
	}
	if (i >= howMany || j >= howMany) {
		return -1;
	}

	r[i].sd = one;
	r[j].sd = two;
	r[i].name = NULL;
	r[j].name = NULL;

	/* let's pair the 2 sockets */
	r[i].companion = j;
	r[j].companion = i;

	return i;
}


static int
AddSocket(	struct host_cookie *cookie,
		Socket sd,
		reflectSocket *r,
		int howMany) {
	int i;
	
	/* sanity check */
	if (sd == NO_SOCKET || r == NULL || cookie == NULL) {
		ERROR("AddSocket: NULL parameter\n");
		return -1;
	}

	/* try to connect to the the host */
	cookie->sd = NO_SOCKET;
	if (!ConnectToHost(cookie, NULL)) {
		return -1;
	}

	/* let's pair the sockets */
	i = PairSockets(sd, cookie->sd, r, howMany);
	if (i == -1) {
		DROP_SOCKET(&sd);
		DROP_SOCKET(&cookie->sd);
		ERROR("AddSocket: reach maximum # of sockets?\n");
		return -1;
	}

	return i;
}


static int
UseSocket(	Socket sd,
		Socket *sList,
		reflectSocket *r,
		int howMany) {
	int i, j;
	
	/* sanity check */
	if (sd == NO_SOCKET || r == NULL || sList == NULL) {
		ERROR("UseSocket: NULL parameter\n");
		return -1;
	}

	/* let's look for a socket in sList */
	for (i = 0; i < howMany; i++) {
		if (sList[i] != NO_SOCKET) {
			break;
		}
	}
	if (i >= howMany) {
		ERROR("UseSocket: no socket from proxy\n");
		DROP_SOCKET(&sd);
		return -1;
	}

	/* all right, let's pair the sockets */
	j = PairSockets(sd, sList[j], r, howMany);
	if (j == -1) {
		ERROR("UseSocket: failed to pair the sockets!\n");
		j = -1;
	} else {
		/* this socket is not anymore available */
		sList[i] = NO_SOCKET;
	}

	return j;
}


static int
ConnectToReflector(	struct host_cookie *cookie,
			Socket *sList,
			int howMany) {
	int i;

	/* sanity check */
	if (cookie == NULL || sList == NULL) {
		ERROR("ConnectToReflector: NULL parameter\n");
		return -1;
	}

	cookie->sd = NO_SOCKET;
	if (!ConnectToHost(cookie, NULL)) {
		ERROR("failed to talk to proxy\n");
		return -1;
	}

	/* let's add the socket into the first available slot */
	for (i = 0; i < howMany; i++) {
		if (sList[i] != NO_SOCKET) {
			break;
		}
	}
	if (i >= howMany) {
		ERROR("ConnectToReflector: out of sockets!\n");
		return -1;
	}
	sList[i] = cookie->sd;

	return i;
}

static int
DeleteSocket(	int index,
		reflectSocket *r) {
 	DROP_SOCKET(&r[index].sd);
	DROP_SOCKET(&r[r[index].companion].sd); 
	r[r[index].companion].companion = NO_SOCKET;
	r[index].companion = NO_SOCKET;
}

static void
DisconnectSocket(Socket sock) {
	int i;

	if (sock == NO_SOCKET) {
		return;
	}

	/* we need to close the matching socket */
	i = FindPairedSocket(sock, myList, HOW_MANY_IN_LIST);
	if (i != -1) {
		/* remove the socket and the companion */
		DeleteSocket(i, myList);
	}

	/* let's see if it's on the other list */
	i = FindSocket(sock, sockList, HOW_MANY_IN_LIST);
	if (i != -1) {
		DROP_SOCKET(&sockList[i]);
	}
}

static void
NewConnection(	Socket ear,
		Socket newSock) {
	int i;

	/* if the ear is for the proxy, put the newSock into the list of
	 * avaiable sockets */
	if (ear == proxyEar) {
		for (i = 0; i < HOW_MANY_IN_LIST; i++) {
			if (sockList[i] == NO_SOCKET) {
				break;
			}
		}
		if (i >= HOW_MANY_IN_LIST) {
			ERROR("NewConnection: out of sockets!\n");
			return;
		}
	}
}


#define SWITCHES "R:v:V:L:"
static void 
usage() {
	printf("\nUsage: nws_reflector [OPTIONS]\n");
	printf("socket reflector for the Network Weather Service\n");
	printf("\nOPTIONS can be:\n");
	printf("\t-R host:port        use #host:port# as our proxy\n");
	printf("\t-L host:port        use #host:port# to connect to\n");
	printf("\t-v level            verbose level (up to 5)\n");
	printf("\t-V                  print version\n");
	printf("Report bugs to <nws@nws.cs.ucsb.edu>.\n\n");
}

int 
main (		int argc, 
		char **argv) {
	Socket sd, ear;
	int i, opt, verbose = 2;
	unsigned short tmp;
	extern char* optarg;
	FILE *logFD, *errFD;
	struct host_cookie rHost, lHost;

	/* defaults */
	logFD = stdout;
	errFD = stderr;
	rHost.name[0] = '\0';
	lHost.name[0] = '\0';
	proxyEar = NO_SOCKET;

	/* let's initialize the list of paired sockets */
	myList = (reflectSocket *)MALLOC(sizeof(reflectSocket)*HOW_MANY_IN_LIST);
	sockList = (Socket *)MALLOC(sizeof(Socket)*HOW_MANY_IN_LIST);
	if (myList == NULL || sockList == NULL) {
		ABORT("out of memory\n");
	}
	for (i = 0; i < HOW_MANY_IN_LIST; i++) {
		myList[i].sd = NO_SOCKET;
		sockList[i] = NO_SOCKET;
	}

	/* no error messages from getopt() */
	opterr = 0;
	
	/* let's parse the command line */
	while((int)(opt = getopt(argc, argv, SWITCHES)) != EOF) {
		switch(opt) {
		case 'L':
			SAFESTRCPY(lHost.name, optarg);
			break;

		case 'R':
			SAFESTRCPY(rHost.name, optarg);
			break;

		case 'V':
			printf("nws_reflector for NWS version %s", VERSION);
#ifdef HAVE_PTHREAD_H
			printf(", with thread support");
#endif
#ifdef WITH_DEBUG
			printf(", with debug support");
#endif
			printf("\n\n");
			exit(0);
			break;

		case 'v':
			verbose = (unsigned short)atol(optarg);
			break;

		case '?':
			if (optopt == 'v') {
				/* using the first level */
				verbose = 2;
				break;
			}
			/* non-recognized options: printing help */

		default:
			usage();
			exit(1);
			break;
		}
	}

	/* set the verbosity */
        DirectDiagnostics(DIAGFATAL, errFD);
        DirectDiagnostics(DIAGDEBUG, DIAGSUPPRESS);
        DirectDiagnostics(DIAGINFO, DIAGSUPPRESS);
        DirectDiagnostics(DIAGLOG, DIAGSUPPRESS);
        DirectDiagnostics(DIAGWARN, DIAGSUPPRESS);
        DirectDiagnostics(DIAGERROR, DIAGSUPPRESS);
        if (verbose > 5) {
                printf("Verbose level too high: set to 5\n");
                verbose = 5;
        }
        switch (verbose) {
        case 5:
                DirectDiagnostics(DIAGDEBUG, logFD);
        case 4:
                DirectDiagnostics(DIAGINFO, logFD);
        case 3:
                DirectDiagnostics(DIAGLOG, logFD);
        case 2:
                DirectDiagnostics(DIAGWARN, logFD);
        case 1:
                DirectDiagnostics(DIAGERROR, errFD);
        }

	if (lHost.name[0] != '\0') {
		if (!Host2Cookie(lHost.name, 8060, &lHost)) {
			ERROR("failed to create cookie\n");
			exit(1);
		}
	}

	/* let's see if we need to connect to the proxy */
	if (rHost.name[0] != '\0') {
		if (!Host2Cookie(rHost.name, 8888, &rHost)) {
			ERROR("failed to create cookie\n");
			exit(1);
		}

		if (ConnectToReflector(&rHost, sockList, HOW_MANY_IN_LIST) == -1) {
			ERROR("failed to talk to proxy\n");
			exit(1);
		}
#if 0
		if (RetrieveObjects(&nsCookie, "hostType=nameserver", &objs)) {
			AddSetToRegistrations(objs, 1, toSync);
			FREE(objs);
		} else {
			ERROR1("main: failed to talk to %s\n", nsCookie.name);
		}
		DisconnectHost(&nsCookie);
#endif
	} else {
		/* we are the proxy we open 2 ears: one for the external
		 * world and one for the proxyee */
		EstablishAnEar(8888, 8888, &proxyEar, &tmp);
		EstablishAnEar(8889, 8889, &ear, &tmp);

		/* we need to be notified about new connection from the
		 * proxy side */
		NotifyOnNewConnection(&NewConnection);
	}

	/* I need to know when sockets get closed */
	if (!NotifyOnDisconnection(&DisconnectSocket)) {
		ERROR("Couldn't register for closing socket\n");
	}

	while (1) {
		if (!IncomingRequest(100, &sd)) {
			/* nothing to do yet */
			continue;
		}

		/* first of all let's see if we have this socket already
		 * paired */
		i = FindPairedSocket(sd, myList, HOW_MANY_IN_LIST);
		if (i == -1) {
			/* socket not yet pair: if I'm local I may need
			 * to create another socket to the reflector */
			if (rHost.name[0] == '\0') {
				i = FindSocket(sd, sockList, HOW_MANY_IN_LIST);
				if (i == -1) {
					ERROR1("Help! socket %d not in any list!\n", sd);
					exit(1);
				}

				/* since the reflector is using this
				 * socket, let me create a new one */
				ConnectToReflector(&rHost, sockList, HOW_MANY_IN_LIST);

				/* now let's pair the socket */
				i = AddSocket(&lHost, sd, myList, HOW_MANY_IN_LIST);
				if (i == -1) {
					ERROR("Failed to find free socket\n");
				}
			} else {
				/* if I'm the remote proxy, let's check
				 * that this socket doesn't belong to the
				 * list of the open socket to the proxyee*/
				if (FindSocket(sd, sockList, HOW_MANY_IN_LIST) == -1) {
					ERROR("socket already in list!\n");
					exit (1);
				}

				/* good, it's a new connection: let's
				 * pick an unused socket to the proxyee
				 * and use it */
				i = UseSocket(sd, sockList, myList, HOW_MANY_IN_LIST);
				if (i < 0) {
					continue;
				}
			}
		}
		
		/* if we cannot read anything, it's likely to be a dropped
		 * connection: clean it up */
		tmp = ReflectSocket(myList[i].sd, myList[myList[i].companion].sd);
		if (tmp == 0) {
			DeleteSocket(i, myList);
		}
	}
}