#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <getopt.h>

#include <libpq-fe.h>

struct Context {
	PGconn *db;
	int count;

	char *buf;
	int buflen;
	int bufpos;
};

static void die(PGconn *db, const char *msg)
{
	if (db)
		fprintf(stderr, "%s: %s\n", msg, PQerrorMessage(db));
	else
		fprintf(stderr, "%s\n", msg);
	exit(1);
}

static void out_flush(struct Context *ctx)
{
	int out;
	if (!ctx->buf)
		return;

	out = write(1, ctx->buf, ctx->bufpos);
	if (out != ctx->bufpos)
		die(NULL, "failed to write file");
	ctx->bufpos = 0;
	ctx->buflen = 0;
	free(ctx->buf);
	ctx->buf = NULL;
}

static void out_char(struct Context *ctx, char c)
{
	if (ctx->bufpos + 1 > ctx->buflen) {
		if (!ctx->buf) {
			ctx->buflen = 16;
			ctx->buf = malloc(ctx->buflen);
			if (!ctx->buf)
				die(NULL, "failed to allocate buffer");
		} else {
			ctx->buflen *= 2;
			ctx->buf = realloc(ctx->buf, ctx->buflen);
			if (!ctx->buf)
				die(NULL, "failed to resize buffer");
		}
	}

	ctx->buf[ctx->bufpos++] = c;
}

static void proc_value(struct Context *ctx, const char *val, int vlen)
{
	int i;
	char c;

	for (i = 0; i < vlen; i++) {
		c = val[i];
		switch (c) {
		case '\\':
			out_char(ctx, '\\');
			out_char(ctx, '\\');
			break;
		case '\t':
			out_char(ctx, '\\');
			out_char(ctx, 't');
			break;
		case '\n':
			out_char(ctx, '\\');
			out_char(ctx, 'n');
			break;
		case '\r':
			out_char(ctx, '\\');
			out_char(ctx, 'r');
			break;
		default:
			out_char(ctx, c);
			break;
		}
	}
}

static void proc_row(struct Context *ctx, PGresult *res, int tup)
{
	int n = PQnfields(res);
	const char *val;
	int i, vlen;

	ctx->count++;

	for (i = 0; i < n; i++) {
		if (i > 0)
			out_char(ctx, '\t');
		if (PQgetisnull(res, tup, i)) {
			out_char(ctx, '\\');
			out_char(ctx, 'N');
			continue;
		}

		vlen = PQgetlength(res, tup, i);
		val = PQgetvalue(res, tup, i);
		proc_value(ctx, val, vlen);
	}
	out_char(ctx, '\n');
	out_flush(ctx);
}

static void exec_query_full(struct Context *ctx, const char *q)
{
	PGconn *db = ctx->db;
	PGresult *r;
	ExecStatusType s;
	int i;

	ctx->count = 0;

	if (!PQsendQuery(db, q))
		die(db, "PQsendQuery");

	/* get next result */
	r = PQgetResult(db);
	s = PQresultStatus(r);
	if (s != PGRES_TUPLES_OK)
		die(db, PQresStatus(s));

	for (i = 0; i < PQntuples(r); i++) {
		proc_row(ctx, r, i);
		ctx->count++;
	}
	PQclear(r);
}

static void exec_query_single_row(struct Context *ctx, const char *q)
{
	PGconn *db = ctx->db;
	PGresult *r;
	ExecStatusType s;

	ctx->count = 0;

	if (!PQsendQuery(db, q))
		die(db, "PQsendQuery");

	if (!PQsetSingleRowMode(db))
		die(NULL, "PQsetSingleRowMode");

	/* loop until all resultset is done */
	while (1) {
		/* get next result */
		r = PQgetResult(db);
		if (!r)
			break;
		s = PQresultStatus(r);
		switch (s) {
			case PGRES_TUPLES_OK:
				//printf("query successful, got %d rows\n", ctx->count);
				ctx->count = 0;
				break;
			case PGRES_SINGLE_TUPLE:
				proc_row(ctx, r, 0);
				break;
			default:
				fprintf(stderr, "result: %s\n", PQresStatus(s));
				exit(1);
				break;
		}

		PQclear(r);
	}
}

#ifdef HAVE_ROWDATA

static void proc_row_zcopy(struct Context *ctx, PGresult *res, PGdataValue *cols)
{
	int n = PQnfields(res);
	const char *val;
	int i, vlen;

	ctx->count++;

	for (i = 0; i < n; i++) {
		if (i > 0)
			out_char(ctx, '\t');
		if (cols[i].len == -1) {
			out_char(ctx, '\\');
			out_char(ctx, 'N');
			continue;
		}

		vlen = cols[i].len;
		val = cols[i].value;
		proc_value(ctx, val, vlen);
	}
	out_char(ctx, '\n');
	out_flush(ctx);
}

static void exec_query_zero_copy(struct Context *ctx, const char *q)
{
	PGconn *db = ctx->db;
	PGresult *r;
	ExecStatusType s;
	PGdataValue *cols;

	ctx->count = 0;

	if (!PQsendQuery(db, q))
		die(db, "PQsendQuery");

	if (!PQsetSingleRowMode(db))
		die(NULL, "PQsetSingleRowMode");

	/* loop until all resultset is done */
	while (PQgetRowData(db, &r, &cols)) {
		proc_row_zcopy(ctx, r, cols);
	}

	/* get final result */
	r = PQgetResult(db);
	s = PQresultStatus(r);
	switch (s) {
	case PGRES_TUPLES_OK:
		//printf("query successful, got %d rows\n", ctx->count);
		ctx->count = 0;
		break;
	default:
		printf("result: %s\n", PQresStatus(s));
		break;
	}
	PQclear(r);
}

#else

static void exec_query_zero_copy(struct Context *ctx, const char *q)
{
	die(NULL, "PQgetRowData() not available");
}

#endif

static const char usage_str[] =
"usage: rowdump [-z|-s|-f] [-d CONNSTR] [-c SQLCOMMAND]\n"
"switches:\n"
"  -f  Load full resultset at once\n"
"  -s  Single-row mode\n"
"  -z  Single-row with direct access (PQgetRowData())\n"
;


static void usage(int err)
{
	printf("%s\n", usage_str);
	exit(err);
}

int main(int argc, char *argv[])
{
	const char *connstr = "dbname=postgres";
	const char *q = "show all";
	PGconn *db;
	struct Context ctx;
	char c;
	int exec_type = 'f';

	while ((c = getopt(argc, argv, "c:d:zsfh")) != -1) {
		switch (c) {
		case 'c':
			q = optarg;
			break;
		case 'd':
			connstr = optarg;
			break;
		case 'z':
		case 's':
		case 'f':
			exec_type = c;
			break;
		case 'h':
			usage(0);
			break;
		default:
			usage(1);
			break;
		}
	}

	db = PQconnectdb(connstr);
	if (!db || PQstatus(db) == CONNECTION_BAD)
		die(db, "connect");

	memset(&ctx, 0, sizeof(ctx));
	ctx.db = db;
	switch (exec_type) {
	case 'z':
		exec_query_zero_copy(&ctx, q);
		break;
	case 'f':
		exec_query_full(&ctx, q);
		break;
	case 's':
		exec_query_single_row(&ctx, q);
		break;
	}

	PQfinish(db);

	return 0;
}