/*
 * Copyright (c) 2006 Teodor Sigaev <teodor@sigaev.ru>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *        notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *        notice, this list of conditions and the following disclaimer in the
 *        documentation and/or other materials provided with the distribution.
 * 3. Neither the name of the author nor the names of any co-contributors
 *        may be used to endorse or promote products derived from this software
 *        without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY CONTRIBUTORS ``AS IS'' AND ANY EXPRESS
 * OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL CONTRIBUTORS BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER
 * IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
 * OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN
 * IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <errno.h>
#include <sys/time.h>

#include "ftsbench.h"

typedef enum RDBMS {
	PostgreSQL = 0,
	MySQL = 1,
	NULLSQL
} RDBMS;

typedef struct RDBMSDesc {
	RDBMS	rdbms;
	char	*shortname;
	char	*longname;
	ftsDB*	(*init)(char *);
} RDBMSDesc;

static RDBMSDesc DBDesc[] = {
	{ PostgreSQL, "pgsql", "PostgreSQL", PGInit }, 
	{ MySQL, 	  "mysql", "MySQL", 	 MYInit },
	{ NULLSQL, 	  NULL,    NULL,    	 NULL   }
};

static void
usage() {
	char buf[1024];
	int i, first=0;

	*buf = '\0';
	for(i=0; DBDesc[i].rdbms != NULLSQL; i++) {
		if ( DBDesc[i].init == NULL )
			continue;
		if ( first != 0 )
			strcat(buf, ", ");
		strcat(buf, DBDesc[i].shortname);
		if ( first == 0 ) 
			strcat(buf, "(default)");
		first++;
	}

	fputs(
		"Copyright (c) 2006 Teodor Sigaev <teodor@sigaev.ru>. All rights reserved.\n"
		"ftsbench - full text search benchmark for RDBMS\n"
		"Initialization of DB:\n"
		"ftsbench -i [-b RDBMS] [-n NUMROW] [-l LEXFILE] [-g GAMMAFILE] [-f FLAGS] [-q] -d DBNAME\n"
		"  -b RDBMS\t- type of DB: ",
		stdout
	);
	fputs( buf, stdout );
	fputs(
		"\n"
		"  -n NUMROW - number of row in table\n"
		"  -l LEXFILE - file with words and its frequents (default gendata/lex)\n"
		"  -g GAMMAFILE - file with doc's length distribution (default gendata/gamma-lens)\n"
		"  -l FLGAS - options for db's schema (see below)\n"
		"  -q - do not print progress message\n",
		stdout
	);
	fputs(
		"Run tests:\n"
		"ftsbench [-b RDBMS] [-c NCLIENTS] [-n NUMQUERY] [-l LEXFILE] [-g GAMMAFILE] [-f FLAGS] [-q] -d DBNAME\n"
		"  -b RDBMS\t- type of DB: ",
		stdout
	);
	fputs( buf, stdout );
	fputs(
		"\n"
		"  -c NCLIENTS - number of clients in parallel\n"
		"  -n NUMQUERY - number of queries per client\n"
		"  -l LEXFILE - file with words and its frequents (default gendata/query-lex)\n"
		"  -g GAMMAFILE - file with doc's length distribution (default gendata/query-lens)\n"
		"  -l FLGAS - options for db's schema (see below)\n"
		"  -q - do not print progress message\n",
		stdout
	);
	fputs(
		"FLAGS are comma-separate list of:\n"
		"  gin  - use GIN index\n"
		"  gist - use GiST index\n"
		"  func - use functional index\n"
		"  and  - AND'ing lexemes in query (default)\n"
		"  or   - OR'ing lexemes in query\n",
		stdout
	);
	fputs(
		"Print SQL-scheme for statistics:\n"
		"ftsbench -S\n",
		stdout
	);
	exit(1);
}

static RDBMS
getRDBMS(char *name) {
	int 	i;

	for(i=0; DBDesc[i].rdbms != NULLSQL; i++) {
		if ( name == NULL ) {
			if ( DBDesc[i].init )
				return DBDesc[i].rdbms; 
		} else if ( strcasecmp(name,DBDesc[i].shortname) == 0 ) {
			if ( DBDesc[i].init == NULL ) {
				fprintf(stderr,"Support of '%s' isn't compiled-in\n", DBDesc[i].longname);
				exit(1);
			}
			return DBDesc[i].rdbms;
		}
	}

	fprintf(stderr,"Can't find a RDBMS\n");
	exit(1);
	
	return NULLSQL;
}

static int
getFLAGS(char *flg) {
	int flags = 0;

	if ( strcasestr(flg,"gist") )
		flags |= FLG_GIST;
	if ( strcasestr(flg,"gin") )
		flags |= FLG_GIN;
	if ( strcasestr(flg,"func") )
		flags |= FLG_FUNC;
	if ( strcasestr(flg,"and") )
		flags |= FLG_AND;
	if ( strcasestr(flg,"or") )
		flags |= FLG_OR;

	if ( (flags & FLG_GIST) && (flags & FLG_GIN) ) {
		fprintf(stderr,"GIN and GiST flags are mutually exclusive\n");
		exit(1);
	}
	if ( (flags & FLG_AND) && (flags & FLG_OR) ) {
		fprintf(stderr,"AND and OR flags are mutually exclusive\n");
		exit(1);
	}

	return flags;
}

static ftsDB **
initConnections(RDBMS rdbms, int n, char *connstr) {
	ftsDB	**dbs = (ftsDB**)malloc(sizeof(ftsDB*) * n);
	int i;

	if (!dbs) {
		fprintf(stderr,"Not enough mwmory\n");
		exit(1);
	}

	for(i=0;i<n;i++) { 
		dbs[i] = DBDesc[rdbms].init(connstr);
		pthread_mutex_init(&dbs[i]->nqueryMutex, NULL);
	}

	return dbs;
}

static double
timediff(struct timeval *begin, struct timeval *end) {
    return ((double)( end->tv_sec - begin->tv_sec )) + ( (double)( end->tv_usec-begin->tv_usec ) ) / 1.0e+6;
}

static double
elapsedtime(struct timeval *begin) {
    struct timeval end;
	gettimeofday(&end,NULL);
	return timediff(begin,&end);
}

static int benchFlags  = 0;
static int benchCount  = 0;
static pthread_cond_t condFinish = PTHREAD_COND_INITIALIZER;
static pthread_mutex_t mutexFinish = PTHREAD_MUTEX_INITIALIZER;
static pthread_mutex_t mutexWordGen = PTHREAD_MUTEX_INITIALIZER;

static void*
execBench(void *in) {
	ftsDB *db = (ftsDB*)in;
	int i;
	char **words;

	for(i=0;i<benchCount;i++) {
		/*
		 * generate_querywords() isn't a thread safe
		 */
		pthread_mutex_lock( &mutexWordGen );
		words = generate_querywords();
		pthread_mutex_unlock( &mutexWordGen );

		db->execQuery(db, words, benchFlags);
		free(words);
	}

	/*
	 * send message about exitting
	 */
    pthread_mutex_lock( &mutexFinish );
	pthread_cond_broadcast( &condFinish );
	pthread_mutex_unlock( &mutexFinish );

	return NULL;	
}

extern char *optarg;

int
main(int argn, char *argv[]) {
	int		initMode = 0;
	int		n = 0, nclients = 1;
	char	*lex = NULL;
	char	*doc = NULL;
	char	*dbname = NULL;
	RDBMS	rdbms = NULLSQL;
	int		flags = 0;
	int i;
	int		quiet = 0, scheme=0;
	StringBuf	b = {NULL,0,0};

	while((i=getopt(argn,argv,"ib:n:l:g:d:c:hf:qS")) != EOF) {
		switch(i) {
			case 'i': initMode = 1; break;
			case 'b': rdbms = getRDBMS(optarg); break;
			case 'n': n=atoi(optarg); break;
			case 'c': nclients=atoi(optarg); break;
			case 'l': lex = strdup(optarg); break;
			case 'g': doc = strdup(optarg); break;
			case 'd': dbname = strdup(optarg); break;
			case 'f': flags = getFLAGS(optarg); break;
			case 'q': quiet = 1; break;
			case 'S': scheme = 1; break;
			case 'h':
			default:
				usage();
		}
	}

	if ( scheme ) {
		printScheme();
		return 0;
	}

	if (rdbms == NULLSQL)
		rdbms = getRDBMS(NULL);

	if ( dbname == NULL || n<0 || nclients<1 )
		usage();

	printf("Running with '%s' RDBMS\n", DBDesc[ rdbms ].longname); 

	if ( initMode ) {
		ftsDB	*db = *initConnections(rdbms, 1, dbname);
		time_t	prev;

		if (!lex)  lex = "gendata/lex";
		if (!doc)  doc = "gendata/gamma-lens";
		finnegan_init(lex, doc);

		db->startCreateScheme(db, flags);
		prev = time(NULL);
		for(i=0;i<n;i++) {
			generate_doc(&b);
			db->InsertRow(db, i+1, b.str);
			if ( !quiet && prev!=time(NULL) ) {
				printf("\r%d(%.02f%%) rows inserted", i, (100.0*i)/n);
				fflush(stdout);
				prev = time(NULL);
			}
		}
		printf("%s%d(100.00%%) rows inserted. Finalyze insertion... ", 
			(quiet) ? "" : "\r", i);
		fflush(stdout);
		db->finishCreateScheme(db);
		printf("done\n");
		db->Close(db);
	} else {
		ftsDB	**dbs = initConnections(rdbms, nclients, dbname);
		pthread_t	*tid = (pthread_t*)malloc( sizeof(pthread_t) * nclients);
		struct 	timeval begin;
		double	elapsed;
		int 	total=0, nres=0;
		struct      timespec  sleepTo = { 0, 0 };

		/*
		 * startup generator
		 */
		if (!lex)  lex = "gendata/query-lex";
		if (!doc)  doc = "gendata/query-lens";
		finnegan_init(lex, doc);

		/*
		 * Initial query
		 */
		if ( !quiet ) {
			printf("\r0(0.00%%) queries proceed");
			fflush(stdout);
		}
		benchFlags = flags;
		benchCount = n;

		gettimeofday(&begin,NULL);

    	pthread_mutex_lock( &mutexFinish );
		for(i=0;i<nclients;i++) {
			if ( pthread_create(tid+i, NULL, execBench, (void*)dbs[i]) != 0 ) {
				fprintf(stderr,"pthread_create failed: %s\n", strerror(errno));
				exit(1);
			}
		}

		for(;;) {
			int res, ntogo = 0;

			total = 0;
			for(i=0;i<nclients;i++) {
				pthread_mutex_lock(&dbs[i]->nqueryMutex);
				total +=dbs[i]->nquery;
				if ( dbs[i]->nquery < n )
					ntogo++;
				pthread_mutex_unlock(&dbs[i]->nqueryMutex);
			}

			if ( ntogo == 0 ) 
				break;

			if ( !quiet ) {
				printf("\r%d(%.02f%%) queries proceed", total, (100.0*(float)total)/(nclients * n));
				fflush(stdout);
			}
			
			sleepTo.tv_sec = time(NULL) + 1;
			res = pthread_cond_timedwait( &condFinish, &mutexFinish, &sleepTo );

			if ( !(res == ETIMEDOUT || res == 0) ) {
				fprintf(stderr,"pthread_cond_timedwait failed: %s\n", strerror(errno));
				exit(1);
			}
		}
		elapsed = elapsedtime(&begin);
		pthread_mutex_unlock( &mutexFinish );

		for(i=0;i<nclients;i++) {
			pthread_join(tid[i], NULL);
			nres += dbs[i]->nres;
			dbs[i]->Close(dbs[i]);
		}

		printf("%s%d(%.02f%%) queries proceed\n", 
			(quiet) ? "" : "\r", total, (100.0*(float)total)/(nclients * n));
		printf("Total number of result: %d\n", nres);
		printf("Total time: %.02f sec, Queries per second: %.02f\n", elapsed, total/elapsed);
		fflush(stdout);
	}

	return 0;
}