#include "trafstats.h"
#include <getopt.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/stat.h>

// Declare threads, database connection and mutexes globally
// TODO: handle various signals to perform graceful shutdown.
//

pthread_t *sniffer_thread=NULL;
pthread_t *storage_thread=NULL;
pthread_t *timer_thread=NULL;
pthread_cond_t timer_gonogo=PTHREAD_COND_INITIALIZER;
pthread_cond_t storage_gonogo=PTHREAD_COND_INITIALIZER;
pthread_cond_t sniffer_gonogo=PTHREAD_COND_INITIALIZER;

// PgDatabase *dbConn=NULL;
// PgDatabase *dbconn=NULL; // TODO: think of a better name
pthread_mutex_t datalock = PTHREAD_MUTEX_INITIALIZER;
pthread_mutex_t timelock = PTHREAD_MUTEX_INITIALIZER;
pthread_mutex_t timergonogolock = PTHREAD_MUTEX_INITIALIZER;
pthread_mutex_t storagegonogolock = PTHREAD_MUTEX_INITIALIZER;
pthread_mutex_t sniffergonogolock = PTHREAD_MUTEX_INITIALIZER;

// static list to be filled.
CTrafStatsList theList;
extern char *optarg;
int verbosity;
char* timestamp = NULL;
bool run=1;
bool castrate;

void *storageThread(void *d);
void *snifferThread(void *d);
void *timerThread(void *d);


void sigHandler(int sig) {
	switch(sig) {
	case SIGINT: // Someone pressed ^C
		//FALLTHRU
	case SIGTERM:
		if(verbosity >= LOG_NOTICE) {
			syslog(LOG_NOTICE,"[main] Signal %d received, exiting",
					sig);
		}
		// Time to leave.
		run=0;
		if(sniffer_thread != NULL) {
			// Pass the signal on to the sniffer thread.
			pthread_kill(*sniffer_thread,sig);
		}
		// once the sniffer thread returns, the storage
		// thread will be warned appropriately by the main
		// thread.
		// 
		break;
	case SIGSEGV:
		if(verbosity >= LOG_INFO) {
			syslog(LOG_INFO,"[main] Something segfaulted.");
			pthread_kill(*sniffer_thread,SIGTERM);
			pthread_kill(*timer_thread,SIGTERM);
			pthread_kill(*storage_thread,SIGTERM);
		}
		exit(sig);
	default:
		if(verbosity >= LOG_WARNING) {
			syslog(LOG_WARNING,"[main] Don't know how to handle signal %d",sig);
		}
	}
}

void usage(const char* program) {
	cerr << "Usage: " << program
		<< " [-i | --interface iface]"
		<< " [-d | --storage-delay delay]"
		<< " [-I | --timer-interval interval]"
		<< " [-D | --database database]"
		<< " [-u | --user user]"
		<< " [-N | --no-daemon ]"
		<< " [-v]"
		<< " [-q]"
		<< endl
		<< "       " << program
		<< " -h" << endl;
}

int main(int argc,char **argv) { 
	char *dev=NULL;
	char *user=NULL;
	char *dbname=NULL;
	char errbuf[PCAP_ERRBUF_SIZE];
	int daemon=true; // Detach into background by default.
	verbosity=4;
	char* pidfile="/var/run/trafstats.pid";

	// A lot of the above and some of the below was shamelessly taken from
	// the libpcap tutorials. Everything that works is theirs, everything
	// that breaks probably mine.
	//
	
	int delay=3600;
	int timerdelay=delay;
	int op=0;
	int pid;
	user="accountant";
	dbname="trafstats";
	char *connstr=NULL;
	static struct option long_options[] = {
		{"interface",1,0,'i'},
		{"device",1,0,'i'},
		{"delay",1,0,'d'},
		{"storage-delay",1,0,'d'},
		{"timer-interval",1,0,'I'},
		{"timer-delay",1,0,'I'},
		{"database",1,0,'D'},
		{"user",1,0,'u'},
		{"verbose",1,0,'V'},
		{"quiet",0,0,'q'},
		{"silent",0,0,'q'},
		{"help",0,0,'h'},
		{"nodaemon",0,0,'N'},
		{"no-daemon",0,0,'N'},
		{"nodetach",0,0,'N'},
		{"no-detach",0,0,'N'},
		{"castrate",0,0,'C'},
		{"ip-only",0,0,'C'},
		{"help",0,0,'h'},
		{0,0,0,0}
	};
	int index=0;

	while ((op = getopt_long(argc,
					argv,
					"V:I:i:d:D:u:CNhvq",
					long_options,
					&index))!=-1)
	{
		switch(op) {

		case 'i':	// Interface
			dev=optarg;
			break;
		case 'd':	// Delay
			delay=atoi(optarg);
			break;
		case 'I':	// Interval (of timer thread)
			timerdelay=atoi(optarg);
			break;
		case 'D':
			dbname=optarg;
			break;
		case 'u':
			user=optarg;
			break;
		case 'v':
			verbosity++;
			break;
		case 'V':
			verbosity=atoi(optarg);
			break;
		case 'q':
			verbosity=0;
			break;
		case 'N':
			daemon=false;
			break;
		case 'C':
			castrate=true;
			break;
		case 'h':
			usage(argv[0]);
			exit(0);
		default:
			usage(argv[0]);
			exit(1);
		}
	}
	if(daemon) {
		// Don't clutter up the console if running as daemon.
		openlog("trafstats",LOG_CONS|LOG_PID,LOG_DAEMON);
	} else {
		openlog("trafstats",LOG_CONS|LOG_PERROR|LOG_PID,LOG_DAEMON);
	}
	if(verbosity >= LOG_NOTICE) {
		syslog(LOG_NOTICE,
			"[main] Using interface %s, delay %d, timer delay %d, database %s, userID %s",
			dev,
			delay,
			timerdelay,
			dbname,
			user);
		if(castrate) {
			syslog(LOG_NOTICE,"Running in castration mode.");
		}
	}

	// Do some sanity checking
	if(delay < timerdelay) {
		if(verbosity >= LOG_WARNING) {
			syslog(LOG_WARNING,"[main] Storage time interval less than timestamp interval; Increasing to %d",timerdelay);
		}
		delay=timerdelay;
	}
	// Grab a device to peek into or exit when we can't find one. 
	if(dev==NULL){
		dev = pcap_lookupdev(errbuf);
		if(dev == NULL) { 
			if(verbosity > LOG_CRIT) {
				syslog(LOG_CRIT,
					"Network device open failed: %s",
					errbuf);
			}
			cerr << errbuf << endl;
			exit(1); 
		}
	}

	// Set up the connection string for the storage and timer threads.
	// 
	unsigned int bufsize=14 + strlen(user) + strlen (dbname);
	connstr=(char*)malloc(bufsize);
	snprintf(connstr,bufsize,"dbname=%s user=%s",dbname,user);
	if(verbosity>=LOG_DEBUG) {
		syslog(LOG_DEBUG,"[main] Database connection string: %s",
				connstr);
	}

	// Initialise the mutexes

	pthread_mutex_init(&datalock,NULL);
	pthread_mutex_init(&timelock,NULL);

	pthread_mutex_init(&timergonogolock,NULL);
	pthread_mutex_init(&storagegonogolock,NULL);
	pthread_mutex_init(&sniffergonogolock,NULL);

	// initialise the timer condition.
	pthread_cond_init(&timer_gonogo,NULL);
	pthread_cond_init(&sniffer_gonogo,NULL);
	pthread_cond_init(&storage_gonogo,NULL);

	// Initialise the threads.
	sniffer_thread=(pthread_t *) malloc(sizeof (pthread_t));
	storage_thread=(pthread_t *) malloc(sizeof (pthread_t));
	timer_thread=(pthread_t *) malloc(sizeof (pthread_t));

	// All preliminary initialising has been done. fork() time!
	if(daemon) {
		if(verbosity >= LOG_INFO) {
			syslog(LOG_INFO,"[main] forking.");
		}
		pid=fork();
		if(pid < 0) { // something bad happened
			switch(errno) {
			case EAGAIN: // Not enough resources.
				syslog(LOG_CRIT, "[main] fork() failed - insufficient resources, exiting.");
				exit(1);
				break; // redundant
			case ENOMEM: // REALLY out of memory.
				syslog(LOG_CRIT,"[main] fork(): Not enough memory to fork, exiting.");
				exit(1);
				break;
			default: // Er?
				syslog(LOG_CRIT,"[main] Unprovided for error during fork(), exiting.");
				exit(1);
			}
		} else if(pid > 0) { // Fork went okay, we're the parent.
			if(verbosity >=LOG_INFO) {
				syslog(LOG_INFO,"[main]: pid is %d",pid);
			}
			// Write pid file.
			umask(026);
			FILE *fd=fopen(pidfile,"w");
			fprintf(fd,"%d",pid);
			exit(0);
		} // Else, we're the child and can get to work.
		if(verbosity >= LOG_DEBUG) {
			syslog(LOG_DEBUG,"[main] Child running.");
		}
	}


	if(verbosity>=LOG_DEBUG) {
		syslog(LOG_DEBUG,"[main]Launching threads.");
	}

	struct timer_options topts = {
		timerdelay,
		connstr
	} ;

	// Launch timer thread.
	pthread_create(timer_thread,NULL,timerThread,(void *)&topts);
	// Wait on timer go/no go
	if(verbosity >= LOG_DEBUG) {
		syslog(LOG_DEBUG,"[main] Calling go/no go.");
		syslog(LOG_DEBUG,"[main] Timer?.");
	}

	// pthread_mutex_lock(&timergonogolock);
	pthread_cond_wait(&timer_gonogo,&timergonogolock);
	// pthread_mutex_unlock(&timergonogolock);

	if(run==0) {
		if(verbosity>=LOG_ERR) {
			syslog(LOG_ERR,"[main] Timer startup responds NO GO, aborting.");
		}
		return 1;
	} // Else
		
	if(verbosity>=LOG_DEBUG) {
		syslog(LOG_DEBUG,"[main] Timer GO.");
	}

	struct storage_options sopts = {
		delay,
		connstr
	} ;
	
	// launch storage thread.
	pthread_create(storage_thread,NULL,storageThread,(void *)&sopts);
	if(verbosity>=LOG_DEBUG) {
		syslog(LOG_DEBUG,"[main] Storage?.");
	}

	//Storage: Go or no go?
	// pthread_mutex_lock(&storagegonogolock);
	pthread_cond_wait(&storage_gonogo,&storagegonogolock);
	// pthread_mutex_unlock(&storagegonogolock);
	if(run==0) {
		if(verbosity>=LOG_ERR) {
			syslog(LOG_ERR,"[main] Storage responds NO GO. Aborting.");
		}
		pthread_join(*storage_thread,NULL);
		if(verbosity>=LOG_DEBUG) {
			syslog(LOG_DEBUG,"[main] Shutting down timer.");
		}
		
		pthread_kill(*timer_thread,SIGUSR1);
		pthread_join(*timer_thread,NULL);
		if(verbosity>=LOG_DEBUG) {
			syslog(LOG_DEBUG,"[main] Timer thread returned. Exiting.");
		}
		return 1;
	} // Else
	
	if(verbosity>=LOG_DEBUG) {
		syslog(LOG_DEBUG,"[main] Storage GO.");
	}

	pthread_create(sniffer_thread,NULL,snifferThread,(void *)dev);

	if(verbosity>=LOG_INFO) { 
		syslog(LOG_INFO,"[main] Setting signal handlers.");
	}

	signal(SIGINT,sigHandler); 
	signal(SIGTERM,sigHandler);
	signal(SIGSEGV,sigHandler);
	/* pthread_join will happily wait on the sniffer thread and still 
	 * respond to signals handled.
	 */

	if(verbosity>=LOG_DEBUG) {
		syslog(LOG_INFO,"[main] Starting wait...");
	}
	pthread_join(*sniffer_thread,NULL);
	if(verbosity >= LOG_DEBUG) {
		syslog(LOG_DEBUG,"[main] sniffer thread returned.");
	}
	run=0;
	if(verbosity >= LOG_DEBUG) {
		syslog(LOG_DEBUG,"[main] Sending kill to storage thread and waiting.");
	}
	pthread_kill(*storage_thread,SIGUSR1);
	pthread_join(*storage_thread,NULL);
	if(verbosity >= LOG_DEBUG) {
		syslog(LOG_DEBUG,"[main] waiting on timer thread.");
	}
	pthread_kill(*timer_thread,SIGUSR1);
	pthread_join(*timer_thread,NULL);
	syslog(LOG_INFO,"[main] All threads returned; exiting.");
	return 0;
}