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

   Vmod-sql 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-sql 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-sql.  If not, see <http://www.gnu.org/licenses/>.
*/
#include "vmod-sql.h"
#include <string.h>
#include <libpq-fe.h>

#define ISSPACE(c) ((c) == ' ' || (c) == '\t')

static char *
chop(char *str)
{
	int len;
  
	if (!str)
		return NULL;
	for (len = strlen(str); len > 0 && ISSPACE(str[len-1]); len--)
		;
	str[len] = 0;
	return str;
}

struct vmod_pgsql_data
{
	PGconn *pgconn;
	PGresult *res;
};

static int
s_pgsql_init(struct vmod_sql_connection *conn)
{
	struct vmod_pgsql_data *dp = calloc (1, sizeof (*dp));
	if (!dp)
		return ENOMEM;
	conn->data = dp;
	return 0;
}

static void
s_pgsql_destroy(struct vmod_sql_connection *conn)
{
	struct vmod_pgsql_data *dp = conn->data;
	free(dp);
	conn->data = NULL;
}

static int
s_pgsql_connect(struct vmod_sql_connection *conn)
{
	struct vmod_pgsql_data *dp = conn->data;

	dp->pgconn = PQsetdbLogin(i_modsql_findparam(conn->param, "server"),
				  i_modsql_findparam(conn->param, "port"),
				  i_modsql_findparam(conn->param, "options"),
				  NULL,
				  i_modsql_findparam(conn->param, "database"),
				  i_modsql_findparam(conn->param, "user"),
				  i_modsql_findparam(conn->param, "password"));

	if (PQstatus(dp->pgconn) == CONNECTION_BAD) {
		i_modsql_error("connection to the database failed: %s",
			   PQerrorMessage(dp->pgconn));
		PQfinish(dp->pgconn);
		dp->pgconn = NULL;
		return 1;
	}
	
	return 0;
}

static int 
s_pgsql_disconnect(struct vmod_sql_connection *conn)
{
	struct vmod_pgsql_data *dp = conn->data;
	PQfinish(dp->pgconn);
	return 0;
}

static int
s_pgsql_query(struct vmod_sql_connection *conn, const char *query)
{
	struct vmod_pgsql_data *dp = conn->data;
	ExecStatusType stat;

	dp->res = PQexec(dp->pgconn, query);
	stat = PQresultStatus(dp->res);

	if (stat == PGRES_TUPLES_OK)
		conn->state = state_result;
	else if (stat != PGRES_COMMAND_OK) {
		i_modsql_error("query failed: %s", PQerrorMessage(dp->pgconn));
		i_modsql_error("the failed query was: %s", query);
		if (dp->res) {
			PQclear(dp->res);
			dp->res = NULL;
		}
		return 1;
	}
	return 0;
}

static int
s_pgsql_free_result(struct vmod_sql_connection *conn)
{
	struct vmod_pgsql_data *dp = conn->data;
	
	PQclear(dp->res);
	dp->res = NULL;

	return 0;
}

static unsigned
s_pgsql_num_fields(struct vmod_sql_connection *conn)
{
	struct vmod_pgsql_data *dp = conn->data;
	if (!dp->res)
		return 0;
	return PQnfields(dp->res);
}

static unsigned
s_pgsql_num_tuples(struct vmod_sql_connection *conn)
{
	struct vmod_pgsql_data *dp = conn->data;
	if (!dp->res)
		return 0;
	return PQntuples(dp->res);
}

static const char *
s_pgsql_get_column(struct vmod_sql_connection *conn, size_t nrow, size_t ncol)
{
	struct vmod_pgsql_data *dp = conn->data;

	if (!dp->res)
		return NULL;
	return chop(PQgetvalue(dp->res, nrow, ncol));
}

static char *
s_pgsql_escape(struct vmod_sql_connection *conn, const char *ustr)
{
	char *str, *q;
	const unsigned char *p;
	size_t len = strlen(ustr);
#define ESCAPABLE_CHAR "\\'\""
  
	for (p = (const unsigned char *) ustr; *p; p++) {
		if (strchr(ESCAPABLE_CHAR, *p))
			len++;
	}

	str = malloc(len + 1);
	if (!str)
		return NULL;

	for (p = (const unsigned char *) ustr, q = str; *p; p++) {
		if (strchr(ESCAPABLE_CHAR, *p))
			*q++ = '\\';
		*q++ = *p;
	}
	*q = 0;
	return str;
}

long
s_pgsql_affected_rows(struct vmod_sql_connection *conn)
{
	struct vmod_pgsql_data *dp = conn->data;
	return strtol(PQcmdTuples(dp->res), NULL, 0);
}
		

struct vmod_sql_backend i_modsql_pgsql_backend = {
	"pgsql",
	s_pgsql_init,
	s_pgsql_destroy,
	s_pgsql_connect,
	s_pgsql_disconnect,
	s_pgsql_escape,
	s_pgsql_query,
	s_pgsql_num_tuples,
	s_pgsql_num_fields,
	s_pgsql_free_result,
	s_pgsql_get_column,
	s_pgsql_affected_rows
};