#include "sqlite3.h"
#include "api.h"
#include "zlib.h"
#include <stdio.h>
#include <string.h>
#include <math.h>
#include <stdlib.h>
#ifdef _WIN32
#include <Windows.h>
#else
#include <unistd.h>
#endif
#include <vector>
#include <string>

#define SQL_TYPE		'SQ3'
#define SQL_MISSING		1.e31

#pragma warning(disable: 4244)

/////////////////////////////////////////////////////////////////////

class zsSqlite {
public:
	zsSqlite() : db(0), ctx(0), func(0), fp(0), nr(0), nr_max(0), nc(0), count(1), callback(0) { strcpy(null,"NULL"); }
	~zsSqlite() { if (db) sqlite3_close(db); if (fp) fclose(fp); }

public:
	sqlite3		*db;
	void		*ctx;
	void		*func;
	FILE		*fp;
	std::string					dstr;
	std::vector<std::string>	dats;
	std::vector<double>			data;
	std::string fname;
	int			nr, nr_max;
	int			nc;
	int			count;
	char		null[8];
	int (*callback)(void*,int,char**,char**);
};

static int file_callback(void* udata, int narg, char **results, char **colnames)
{
	zsSqlite *o = (zsSqlite*)udata;
	if (o->nr_max > 0 && o->count >= o->nr_max) return 0;
	if (!o->fp) return 0;

	if (results[0]) o->dstr=results[0];
	for (int i = 1; i < narg; i++) {
		o->dstr+=",";
		if (results[i] && strlen(results[i]) > 0) {
			o->dstr+=results[i];
		}
		else {
			o->dstr+=o->null;
		}
	}
	fprintf(o->fp,"%s\n",o->dstr.c_str());
	return 0;
}

static int pack_callback(void* udata, int narg, char **results, char **colnames)
{
	zsSqlite *o = (zsSqlite*)udata;
	if (o->nr_max > 0 && o->count >= o->nr_max) return 0;

	if (results[0]) o->dstr+=results[0];
	for (int i = 1; i < narg; i++) {
		o->dstr+=",";
		if (results[i] && strlen(results[i]) > 0) {
			o->dstr+=results[i];
		}
		else {
			o->dstr+=o->null;
		}
	}
	o->dstr+="\n";
	return 0;
}

static int default_callback(void* udata, int narg, char **results, char **colnames)
{
	zsSqlite *o = (zsSqlite*)udata;
	if (o->nr_max > 0 && o->count >= o->nr_max) return 0;
	if (o->count == 0) {
		for (int i = 0; i < narg; i++) printf("|%s", colnames[i]);
		printf("|\n");
	}
	o->count++;
	for (int i = 0; i < narg; i++) {
		if (results[i] && strlen(results[i]) > 0) {
			printf("|%s", results[i]);
		}
		else {
			printf("|%s", o->null);
		}
	}
	printf("|\n");
	fflush(stdout);
	return 0;
}

static int zs_callback(void* udata, int narg, char **results, char **colnames)
{
	zsSqlite *o = (zsSqlite*)udata;
	if (o->nr_max > 0 && o->count >= o->nr_max) return 0;
	void *p[4];
	if (o->count == 0) {
		for (int i = 0; i < narg; i++) {
			p[0] = api_create_integer(o->ctx, 0);
			p[1] = api_create_integer(o->ctx, i+1);
			p[2] = api_create_string(o->ctx, (char*)colnames[i]);
			p[3] = api_create_integer(o->ctx, narg);
			api_call_func(o->ctx, o->func, 4, p);
		}
	}
	o->count++;
	for (int i = 0; i < narg; i++) {
		p[0] = api_create_integer(o->ctx, o->count);
		p[1] = api_create_integer(o->ctx, i+1);
		if (results[i] && strlen(results[i]) > 0) {
			p[2] = api_create_string(o->ctx, results[i]);
		}
		else {
			p[2] = api_create_string(o->ctx, o->null);
		}
		p[3] = api_create_integer(o->ctx, narg);
		api_call_func(o->ctx, o->func, 4, p);
	}
	return 0;
}

static int string_callback(void* udata, int narg, char **results, char **colnames)
{
	zsSqlite *o = (zsSqlite*)udata;
	if (o->nr_max > 0 && o->nr >= o->nr_max) return 0;
	o->nc = narg;
	for (int i = 0; i < narg; i++) {
		if (results[i] && strlen(results[i]) > 0) {
			o->dats.push_back(results[i]);
		}
		else {
			o->dats.push_back(o->null);
		}
	}
	o->nr++;
	return 0;
}

static int real_callback(void* udata, int narg, char **results, char **colnames)
{
	zsSqlite *o = (zsSqlite*)udata;
	if (o->nr_max > 0 && o->nr >= o->nr_max) return 0;
	o->nc = narg;
	for (int i = 0; i < narg; i++) {
		if (results[i] && strlen(results[i]) > 0) {
			o->data.push_back(atof(results[i]));
		}
		else {
			o->data.push_back(SQL_MISSING);
		}
	}
	o->nr++;
	return 0;
}

static void sql_destroy(void* ptr)
{
	zsSqlite *o = (zsSqlite*)ptr;
	delete o;
    ptr = 0;
}

/////////////////////////////////////////////////////////////////////

static void* sql_create(void *ctx, int nargs, void** args)
{
	if (nargs < 1) api_input_error(ctx);
	zsSqlite *o = new zsSqlite;
	o->ctx = ctx;
	o->fname = api_get_string(ctx, args[0]);
	int status = sqlite3_open(o->fname.c_str(), &o->db);
	int count = 0;
	while (status == SQLITE_BUSY && count < 60) {
#ifdef _WIN32
		Sleep(1000);
#else
		sleep(1);
#endif
		status = sqlite3_open(o->fname.c_str(), &o->db);
		count++;
	}
	if (status == SQLITE_OK) {
		o->callback = &default_callback;
		sqlite3_busy_timeout(o->db,1000);
		return api_create_user(ctx, o, 0, sql_destroy, SQL_TYPE);
	}
	delete o;
	return 0;
}

static void* sql_reopen(void *ctx, int nargs, void** args)
{
	if (nargs < 1) api_input_error(ctx);
	zsSqlite *o = (zsSqlite*)api_get_user(ctx, args[0], SQL_TYPE);
	sqlite3_close(o->db);
	int status = sqlite3_open(o->fname.c_str(), &o->db);
	o->ctx = ctx;
	return api_create_integer(ctx, status==SQLITE_OK);
}


static void* sql_version(void *ctx, int nargs, void** args)
{
	if (nargs < 1) api_input_error(ctx);
	zsSqlite *o = (zsSqlite*)api_get_user(ctx, args[0], SQL_TYPE);
	o->ctx = ctx;
	return api_create_string(ctx, (char*)sqlite3_libversion());
}

static void* sql_error(void *ctx, int nargs, void** args)
{
	if (nargs < 1) api_input_error(ctx);
	zsSqlite *o = (zsSqlite*)api_get_user(ctx, args[0], SQL_TYPE);
	o->ctx = ctx;
	return api_create_string(ctx, sqlite3_errmsg(o->db));
}

static void* sql_callback(void *ctx, int nargs, void** args)
{
	if (nargs < 2) api_input_error(ctx);
	zsSqlite *o = (zsSqlite*)api_get_user(ctx, args[0], SQL_TYPE);
	o->ctx = ctx;
	o->func = api_get_func(ctx, api_get_string(ctx, args[1]));
	o->callback = &zs_callback;
	return 0;
}

static void* sql_nrow(void *ctx, int nargs, void** args)
{
	if (nargs < 1) api_input_error(ctx);
	zsSqlite *o = (zsSqlite*)api_get_user(ctx, args[0], SQL_TYPE);
	o->ctx = ctx;
	return api_create_integer(ctx, o->nr);
}

static void* sql_ncol(void *ctx, int nargs, void** args)
{
	if (nargs < 1) api_input_error(ctx);
	zsSqlite *o = (zsSqlite*)api_get_user(ctx, args[0], SQL_TYPE);
	o->ctx = ctx;
	return api_create_integer(ctx, o->nc);
}

static void* sql_query(void *ctx, int nargs, void** args)
{
	if (nargs < 2) api_input_error(ctx);
	zsSqlite *o = (zsSqlite*)api_get_user(ctx, args[0], SQL_TYPE);
	o->ctx = ctx;
	const char *s = api_get_string(ctx, args[1]);
	o->count = 0;
	int (*callback)(void*,int,char**,char**) = o->callback;
	if (nargs > 2) {
		o->nr = 0;
		o->nc = 0;
		o->dats.clear();
		o->data.clear();
		o->dstr="";
		if (api_is_string(args[2])) {
			const char *fname = api_get_string(ctx, args[2]);
			if (o->fp) fclose(o->fp);
			o->fp = fopen(fname, "w");
			if (!o->fp) api_runtime_error2(ctx, "Cannot open ", fname);
			callback = &file_callback;
		}
		else {
			int save = api_get_integer(ctx, args[2]);
			if (save == 1) {
				callback = &string_callback;
			}
			else if (save == 2) {
				callback = &real_callback;
			}
			else {
				callback = &pack_callback;
			}
		}
	}
	char *msg = 0;
	int status = sqlite3_exec(o->db, s, callback, o, &msg);
	int count = 0;
	while (status == SQLITE_BUSY && count < 60) {
		if (msg) {
			sqlite3_free(msg);
			msg = 0;
		}
#ifdef _WIN32
		Sleep(1000);
#else
		sleep(1);
#endif
		status = sqlite3_exec(o->db, s, callback, o, &msg);
		count++;
	}
	if (o->fp) {
		fclose(o->fp);
		o->fp = 0;
	}
	if (status != SQLITE_OK) {
		sqlite3_free(msg);
		return api_create_integer(ctx, 0);
	}
	return api_create_integer(ctx, 1);
}

static void* sql_get(void *ctx, int nargs, void** args)
{
	if (nargs < 2) api_input_error(ctx);
	zsSqlite *o = (zsSqlite*)api_get_user(ctx, args[0], SQL_TYPE);
	int i = api_get_integer(ctx, args[1]);
	if (nargs == 2) {
		// get a row
		if (i < 0 || i >= o->nr) api_runtime_error(ctx, "Bad row index!");
		std::string s;
		int k = i*o->nc;
		if (o->dats.size() > 0) {
			s = o->dats[k];
			for (int j=1; j<o->nc; j++) {
				s += ", ";
				s += o->dats[k+j];
			}
		}
		else {
			char c[128];
			sprintf(c,"%g", o->data[k]);
			s = c;
			for (int j=1; j<o->nc; j++) {
				sprintf(c,", %g", o->data[k+j]);
				s += c;
			}
		}
		return api_create_string(ctx, s.c_str());
	}
	// get a cell
	int j = api_get_integer(ctx, args[2]);
	int k = i*o->nc+j;
	if (o->dats.size() > 0) {
		if (k < 0 || k >= o->dats.size()) api_runtime_error(ctx, "Bad string index!");
		return api_create_string(ctx, o->dats[k].c_str());
	}
	if (k < 0 || k >= o->data.size()) api_runtime_error(ctx, "Bad real index!");
	return api_create_real(ctx, o->data[k]);
}

static void* sql_ptr(void *ctx, int nargs, void** args)
{
	if (nargs < 1) api_input_error(ctx);
	zsSqlite *o = (zsSqlite*)api_get_user(ctx, args[0], SQL_TYPE);
	if (o->data.size() == 0) api_runtime_error(ctx, "Query data not collected as real!");
	void *arr = api_create_array(ctx,3);
	api_set_array_object(ctx, arr, "0", api_create_user(0,&o->data[0],0,0,0));
	api_set_array_object(ctx, arr, "1", api_create_integer(0,o->data.size()));
	api_set_array_object(ctx, arr, "2", api_create_integer(0,sizeof(double)));
	return arr;
}

static void* sql_missing(void *ctx, int nargs, void** args)
{
	if (nargs < 1) api_input_error(ctx);
	zsSqlite *o = (zsSqlite*)api_get_user(ctx, args[0], SQL_TYPE);
	return api_create_real(ctx, SQL_MISSING);
}

static void* sql_limit(void *ctx, int nargs, void** args)
{
	if (nargs < 2) api_input_error(ctx);
	zsSqlite *o = (zsSqlite*)api_get_user(ctx, args[0], SQL_TYPE);
	o->nr_max = api_get_integer(ctx, args[1]);
	return 0;
}

static void* sql_pack(void *ctx, int nargs, void** args)
{
	if (nargs < 1) api_input_error(ctx);
	zsSqlite *o = (zsSqlite*)api_get_user(ctx, args[0], SQL_TYPE);
	int zip = 0;
	if (nargs > 1) zip = api_get_integer(ctx, args[1]);
	if (!zip) {
		return api_create_string(ctx,o->dstr.c_str());
	}

	z_stream strm;
	memset(&strm,0,sizeof(strm));
	strm.zalloc=Z_NULL;
	strm.zfree=Z_NULL;
	strm.opaque=Z_NULL;
	strm.next_in=(Bytef*)o->dstr.c_str();
	strm.avail_in=o->dstr.length();

	if (deflateInit2(&strm,Z_DEFAULT_COMPRESSION,Z_DEFLATED,15|16,8,Z_DEFAULT_STRATEGY)!=Z_OK) {
		fprintf(stderr,"deflateInit2() failed!\n");
		return 0;
	}
	int len=deflateBound(&strm,o->dstr.length())+32;

	void *dst=malloc(len);
	if (!dst) {
		fprintf(stderr,"malloc() failed!\n");
		deflateEnd(&strm);
		return 0;
	}
	strm.next_out=(Bytef*)dst;
	strm.avail_out=len;
	if (deflate(&strm,Z_FINISH)==Z_STREAM_ERROR) {
		fprintf(stderr,"deflate() failed!\n");
		deflateEnd(&strm);
		free(dst);
		return 0;
	}
	if (strm.avail_in!=0) {
		fprintf(stderr,"deflate() incomplete!\n");
		deflateEnd(&strm);
		free(dst);
		return 0;
	}

	len-=strm.avail_out;

	void *arr = api_create_array(ctx,2);
	api_set_array_object(ctx, arr, "0", api_create_user(0,dst,0,free,0));
	api_set_array_object(ctx, arr, "1", api_create_integer(0,len));
	return arr;
}

static void* sql_null(void *ctx, int nargs, void** args)
{
	if (nargs < 2) api_input_error(ctx);
	zsSqlite *o = (zsSqlite*)api_get_user(ctx, args[0], SQL_TYPE);
	if (api_get_integer(ctx, args[1])) {
		strcpy(o->null,"NULL");
	}
	else {
		o->null[0]=0;
	}
	return 0;
}


/////////////////////////////////////////////////////////////////////

class sqliteRegPrimitive
{
public:
	sqliteRegPrimitive()
	{
		api_add_primitive("sqlite",		0,			sql_create);
		api_add_primitive("reset",		SQL_TYPE,	sql_reopen);
		api_add_primitive("version",	SQL_TYPE,	sql_version);
		api_add_primitive("error",		SQL_TYPE,	sql_error);
		api_add_primitive("callback",	SQL_TYPE,	sql_callback);
		api_add_primitive("nrow",		SQL_TYPE,	sql_nrow);
		api_add_primitive("ncol",		SQL_TYPE,	sql_ncol);
		api_add_primitive("query",		SQL_TYPE,	sql_query);
		api_add_primitive("__get",		SQL_TYPE,	sql_get);
		api_add_primitive("ptr",		SQL_TYPE,	sql_ptr);
		api_add_primitive("missing",	SQL_TYPE,	sql_missing);
		api_add_primitive("limit",		SQL_TYPE,	sql_limit);
		api_add_primitive("pack",		SQL_TYPE,	sql_pack);
		api_add_primitive("NULL",		SQL_TYPE,	sql_null);
	}
};

static sqliteRegPrimitive sql_register;