#include "api.h"
#include "buffer.h"
#include "libpq-fe.h"
#include <stdio.h>
#include <string>

#define PG_TYPE		'PQ'

#pragma warning(disable: 4244)

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

class pgSQL {
public:
	pgSQL() : conn(0), res(0) { }

	~pgSQL() { if (conn) PQfinish(conn); if (res) PQclear(res); }

	PGconn *conn;

	PGresult *res;
};


void pg_destroy(void* ptr)
{
    delete (pgSQL*)ptr;
    ptr = 0;
}

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

void* pg_create(void *ctx, int nargs, void** args)
{
	if (nargs < 1) api_input_error(ctx);
	const char *s = api_get_string(ctx, args[0]);
	PGconn *conn = PQconnectdb(s);
	if (PQstatus(conn) != CONNECTION_OK) {
		api_runtime_error(ctx, PQerrorMessage(conn));
	}
	pgSQL *o = new pgSQL;
	o->conn = conn;
	return api_create_user(ctx, o, 0, pg_destroy, PG_TYPE);
}

void* pg_version(void *ctx, int nargs, void** args)
{
	if (nargs < 1) api_input_error(ctx);
	pgSQL *o = (pgSQL*)api_get_user(ctx, args[0], PG_TYPE);
	return api_create_integer(ctx, PQserverVersion(o->conn));
}

void* pg_connect(void *ctx, int nargs, void** args)
{
	if (nargs < 2) api_input_error(ctx);
	pgSQL *o = (pgSQL*)api_get_user(ctx, args[0], PG_TYPE);
	const char *s = api_get_string(ctx, args[1]);
	if (o->conn) {
		PQfinish(o->conn);
		o->conn = 0;
	}
	o->conn = PQconnectdb(s);
	if (PQstatus(o->conn) != CONNECTION_OK) {
		api_runtime_error(ctx, PQerrorMessage(o->conn));
	}
	return 0;
}

void* pg_query(void *ctx, int nargs, void** args)
{
	if (nargs < 2) api_input_error(ctx);
	pgSQL *o = (pgSQL*)api_get_user(ctx, args[0], PG_TYPE);
	if (o->res) {
		PQclear(o->res);
		o->res = 0;
	}
	o->res = PQexec(o->conn, api_get_string(ctx, args[1]));
	ExecStatusType status = PQresultStatus(o->res);
	if (status != PGRES_TUPLES_OK) {
		return api_create_string(ctx, PQresStatus(status));
	}
	void *arr = api_create_array(ctx, 2);
	api_set_array_object(ctx, arr, "0", api_create_integer(0, PQntuples(o->res)));
	api_set_array_object(ctx, arr, "1", api_create_integer(0, PQnfields(o->res)));
	return arr;
}

void* pg_fetch(void *ctx, int nargs, void** args)
{
	if (nargs < 2) api_input_error(ctx);
	pgSQL *o = (pgSQL*)api_get_user(ctx, args[0], PG_TYPE);
	if (!o->res) return 0;
	int nrow = PQntuples(o->res);
	int ncol = PQnfields(o->res);
	if (nrow < 1 || ncol < 1) return 0;
	int k = api_get_integer(ctx, args[1]);
	if (k < 0 || k >= nrow) return 0;
	if (ncol == 1) {
		return api_create_string(ctx, PQgetvalue(o->res, k, 0));
	}
	void *arr = api_create_array(ctx, ncol);
	for (int i = 0; i < ncol; i++) {
		api_set_array_object2(ctx, arr, i, api_create_string(0, PQgetvalue(o->res, k, i)));
	}
	return arr;
}

double cal2jul(const char *s)
{
	char c[5];
	strncpy(c, s, 4);
	c[4] = 0;
	int yy = atoi(c);
	strncpy(c, s+5, 2);
	c[2] = 0;
	int mm = atoi(c);
	strncpy(c, s+8, 2);
	int dd = atoi(c);
	yy -= mm < 3;
	mm += 12 * (mm < 3);
	int a = 365.25 * (yy + 4716),
		b = 30.6001 * (mm + 1);
	double d = a + b + dd - 2416556;
	if (strlen(s) > 18) {
		strncpy(c, s+11, 2);
		d += atof(c)/24.0;
		strncpy(c, s+14, 2);
		d += atof(c)/1440.0;
		strncpy(c, s+17, 2);
		d += atof(c)/86400.0;
	}
	return d;
}

void* pg_data(void *ctx, int nargs, void** args)
{
	if (nargs < 3) api_input_error(ctx);
	pgSQL *o = (pgSQL*)api_get_user(ctx, args[0], PG_TYPE);
	double *d = (double*)api_get_ptr(ctx, args[1]);
	int nd = api_get_integer(ctx, args[2]);
	bool flag = false;
	if (nargs > 3) flag = api_get_integer(ctx, args[3])!=0;
	if (!o->res) return api_create_integer(ctx, 0);
	int nrow = PQntuples(o->res);
	int ncol = PQnfields(o->res);
	if (nrow < 1 || ncol < 1) return api_create_integer(ctx, 0);
	int i, j, k=0;
	for (i = 0; i < nrow; i++) {
		for (j = 0; j < ncol; j++) {
			if (k >= nd) return api_create_integer(ctx, k);
			if (flag && j==0)
				d[k++] = cal2jul(PQgetvalue(o->res, i, j));
			else
				d[k++] = atof(PQgetvalue(o->res, i, j));
		}
	}
	return api_create_integer(ctx, k);;
}

void* pg_print(void *ctx, int nargs, void** args)
{
	if (nargs < 1) api_input_error(ctx);
	pgSQL *o = (pgSQL*)api_get_user(ctx, args[0], PG_TYPE);
	if (!o->res) return 0;
	PQprintOpt po = { 1, 1, 0, 0, 1, 0, ",", 0, 0, 0 };
	PQprint(stdout, o->res, &po);
	return 0;
}

void* pg_escape(void *ctx, int nargs, void** args)
{
	if (nargs < 2) api_input_error(ctx);
	pgSQL *o = (pgSQL*)api_get_user(ctx, args[0], PG_TYPE);
	if (api_is_string(args[1])) {
		const char *s = api_get_string(ctx, args[1]);
		int flag = 0;
		if (nargs > 2) flag = api_get_integer(ctx, args[2]);
		if (flag != 0) {
			size_t n = strlen(s);
			unsigned char *p = PQunescapeBytea((unsigned char*)s, &n);
			if (!p) return 0;
			return api_create_user(ctx, p, 0, PQfreemem, 0);
		}
		int n = strlen(s);
		int error;
		zsBuffer b(n*3);
		n = PQescapeStringConn(o->conn, b.u.pchar, s, n, &error);
		if (error != 0) return 0;
		return api_create_string(ctx, b.u.pchar);
	}
	void *ptr = api_get_ptr(ctx, args[1]);
	size_t n = api_get_integer(ctx, args[2]);
	unsigned char *p = PQescapeByteaConn(o->conn, (unsigned char*)ptr, n, &n);
	if (!p) return 0;
	void *ret = api_create_string(ctx, (char*)p);
	PQfreemem(p);
	return ret;
}

void* pg_unescape(void *ctx, int nargs, void** args)
{
	if (nargs < 2) api_input_error(ctx);
	pgSQL *o = (pgSQL*)api_get_user(ctx, args[0], PG_TYPE);
	const char *s = api_get_string(ctx, args[1]);
	size_t n = strlen(s);
	unsigned char *p = PQunescapeBytea((const unsigned char*)s, &n);
	if (!p) return 0;
	void *arr = api_create_array(ctx, 2);
	api_set_array_object(ctx, arr, "0", api_create_user(0, p, 0, PQfreemem, 0));
	api_set_array_object(ctx, arr, "1", api_create_integer(0, n));
	return arr;
}

void* pg_reset(void *ctx, int nargs, void** args)
{
	if (nargs < 1) api_input_error(ctx);
	pgSQL *o = (pgSQL*)api_get_user(ctx, args[0], PG_TYPE);
	PQreset(o->conn);
	return 0;
}

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

class zsRegPrimitive
{
public:
	zsRegPrimitive()
	{
		api_add_primitive("pgsql",		0,			pg_create);
		api_add_primitive("version",	PG_TYPE,	pg_version);
		api_add_primitive("connect",	PG_TYPE,	pg_connect);
		api_add_primitive("query",		PG_TYPE,	pg_query);
		api_add_primitive("fetch",		PG_TYPE,	pg_fetch);
		api_add_primitive("data",		PG_TYPE,	pg_data);
		api_add_primitive("print",		PG_TYPE,	pg_print);
		api_add_primitive("escate",		PG_TYPE,	pg_escape);
		api_add_primitive("unescate",	PG_TYPE,	pg_unescape);
		api_add_primitive("reset",		PG_TYPE,	pg_reset);
	}

	~zsRegPrimitive() { }
};

zsRegPrimitive pg_register;