/* This file is part of vmod-dbrw
   Copyright (C) 2013-2014 Sergey Poznyakoff

   Vmod-dbrw is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; either version 3, or (at your option)
   any later version.

   Vmod-dbrw is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with vmod-dbrw.  If not, see <http://www.gnu.org/licenses/>.
*/

#include "dbrw.h"
#include "wordsplit.h"
#include <unistd.h>
#include <stdarg.h>
#include <stdio.h>
#include <ctype.h>

const char *progname;

void
dbrw_debug(const char *fmt, ...)
{
	va_list ap;
	fprintf(stderr, "%s: DEBUG: ", progname);
	va_start(ap, fmt);
	vfprintf(stderr, fmt, ap);
	va_end(ap);
	fputc('\n', stderr);
}

void
dbrw_error(const char *fmt, ...)
{
	va_list ap;
	fprintf(stderr, "%s: ", progname);
	va_start(ap, fmt);
	vfprintf(stderr, fmt, ap);
	va_end(ap);
	fputc('\n', stderr);
}	

void
usage(int c)
{
	FILE *fp = c ? stderr : stdout;
	fprintf(fp, "usage: %s [-ch] BACKEND CONN\n", progname);
	fprintf(fp, "Initializes the test database\n\n");
	fprintf(fp, "BACKEND is mysql or pgsql\n");
	fprintf(fp, "CONN is the vmod-dbrw connection parameter string\n");
	fprintf(fp, "(see second argument to dbrw.config in vmod-dbrw(1))\n");
	fprintf(fp, "\nOPTIONS:\n\n");
	fprintf(fp, "   -c      create the database if it does not exist\n");
	fprintf(fp, "   -h      print this help summary\n");
	exit(c);
}


int line;
char *qbuf;
size_t qsize;
size_t qlen;

int
iscomm(const char *p)
{
	for (; *p; p++) {
		if (*p == '-' && p[1] == '-')
			return 1;
		if (!isspace(*p))
			return 0;
	}
	return 1;
}

char *
getquery(FILE *fp)
{
	size_t len;
	size_t off;
	
	qlen = off = 0;
	while (1) {
		if (qlen + 1 >= qsize) {
			if (qsize == 0)
				qsize = 80;
			else
				qsize *= 2;
			qbuf = realloc(qbuf, qsize);
			if (!qbuf) {
				dbrw_error("out of memory");
				abort();
			}
		}
				
			
		++line;
		if (!fgets(qbuf + qlen, qsize - qlen, fp))
			break;
		len = strlen(qbuf + qlen);
		qlen += len;
		len = qlen;
		if (qbuf[len - 1] != '\n')
			continue;
		while (len > 0 && isspace(qbuf[len - 1]))
			--len;
		if (iscomm(qbuf + off)) {
			qlen = off;
			continue;
		} else
			off = qlen;
		if (qbuf[len - 1] == ';') {
			qlen = len - 1;
			break;
		}
	}
	if (qlen == 0)
		return NULL;
	qbuf[qlen] = 0;
	return qbuf;
}

void
trycreate(struct dbrw_connection *conn, const char *dbname)
{
	char q[1024];
	if (strcmp(conn->conf->backend->name, "mysql") == 0) {
		snprintf(q, sizeof q, "CREATE DATABASE IF NOT EXISTS %s",
			 dbname);
		if (sql_query(conn, q)) {
			dbrw_error("query failed: %s", q);
			exit(4);
		}
	} else {
		snprintf(q, sizeof q,
			 "SELECT datname FROM pg_database WHERE datname='%s'",
			 dbname);
		if (sql_query(conn, q)) {
			dbrw_error("query failed: %s", q);
			exit(4);
		}
		if (sql_num_tuples(conn) == 0) {
			snprintf(q, sizeof q, "CREATE DATABASE %s",
				 dbname);
			if (sql_query(conn, q)) {
				dbrw_error("query failed: %s", q);
				exit(4);
			}
		}
	}
}
		   
int
main(int argc, char **argv)
{
	char *s;
	struct wordsplit ws;
	struct dbrw_config cfg;
	struct dbrw_connection conn;
	char buf[1024];
	char *p;
	char *dbname = NULL;
	int create = 0;
	int c;
	int i;
	
	progname = argv[0];

	while ((c = getopt(argc, argv, "ch")) != EOF) {
		switch (c) {
		case 'c':
			create = 1;
			break;

		case 'h':
			usage(0);
			break;

		default:
			return 1;
		}
	}

	argc -= optind;
	argv += optind;
	
	if (argc != 2)
		usage(1);

	memset(&cfg, 0, sizeof(cfg));
	cfg.backend = dbrw_backend_select(argv[0]);
	if (!cfg.backend) {
		dbrw_error("unsupported backend: %s", argv[0]);
		return 1;
	}

	ws.ws_delim = ";";
	wordsplit(argv[1], &ws, 
		  WRDSF_NOVAR | WRDSF_NOCMD | WRDSF_QUOTE |
		  WRDSF_CESCAPES |
		  WRDSF_DELIM |
		  WRDSF_SHOWERR|WRDSF_ENOMEMABRT);

	cfg.param = ws.ws_wordv;
	ws.ws_wordv = NULL;
	ws.ws_wordc = 0;
	wordsplit_free(&ws);
	if (create) {
		for (i = 0; cfg.param[i]; i++)
			if (strncmp(cfg.param[i], "database=", 9) == 0)
				break;
		if (!cfg.param[i]) {
			dbrw_error("no database specified");
			return 1;
		}
		dbname = cfg.param[i];
		for (; cfg.param[i+1]; i++)
			cfg.param[i] = cfg.param[i + 1];
		if (strcmp(cfg.backend->name, "pgsql") == 0) {
			cfg.param[i] = "database=postgres";
			cfg.param[i + 1] = NULL;
		} else
			cfg.param[i] = NULL;
	}
	
	s = findparam(cfg.param, "debug");
	cfg.debug_level = s ? atoi(s) : 0;

	/* Initialize the connection */
	memset(&conn, 0, sizeof(conn));
	conn.state = state_init;
	conn.conf = &cfg;

	if (sql_init(&conn))
		return 2;

	if (sql_connect(&conn) || conn.state != state_connected)
		return 2;

	if (dbname) {
		trycreate(&conn, dbname + 9);
		cfg.param[i] = dbname;
		cfg.param[i+1] = NULL;
		sql_disconnect(&conn);
		if (sql_connect(&conn) || conn.state != state_connected)
			return 2;
	}

	line = 0;
	while ((p = getquery(stdin))) {
		if (sql_query(&conn, p)) {
			dbrw_error("input line %d: query failed: %s",
				   line, p);
			return 4;
		}
	}
	
	sql_disconnect(&conn);
	sql_destroy(&conn);

	return 0;
}