/* This file is part of vmod-sql
   Copyright (C) 2013-2014 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 <mysql/mysql.h>
#include <mysql/errmsg.h>
#include <mysql/mysqld_error.h>

struct vmod_mysql_data
{
	MYSQL *mysql;
	MYSQL_RES *result;
	char *buffer;
	size_t bufsize;
};
  

static void
check_errno(struct vmod_sql_connection *conn)
{
	struct vmod_mysql_data *mp = conn->data;

	switch (mysql_errno(mp->mysql)) {
	case CR_SERVER_GONE_ERROR:
	case CR_SERVER_LOST:
	case ER_SERVER_SHUTDOWN:
	case ER_ABORTING_CONNECTION:
		i_modsql_error("query failed: %s", mysql_error(mp->mysql));
		i_modsql_disconnect(conn);
		if (conn->state == state_error) {
			conn->state = state_disabled;
			i_modsql_error("disabling MySQL connection");
		}
		break;
	case ER_UNKNOWN_COM_ERROR:
	case ER_ACCESS_DENIED_ERROR:
	case ER_BAD_DB_ERROR:
	case ER_WRONG_DB_NAME:
	case ER_BAD_FIELD_ERROR:
	case ER_BAD_HOST_ERROR:
	case ER_BAD_TABLE_ERROR:
	case ER_WRONG_FIELD_SPEC:
	case ER_PARSE_ERROR:
	case ER_EMPTY_QUERY:
	case ER_FIELD_SPECIFIED_TWICE:
	case ER_NO_SUCH_TABLE:
	case ER_NOT_ALLOWED_COMMAND:
		i_modsql_error("query failed: %s", mysql_error(mp->mysql));
	}
}

/* ************************************************************************* */
/* Interface routines */

static int
s_mysql_init(struct vmod_sql_connection *conn)
{
	struct vmod_mysql_data *mp = calloc(1, sizeof (*mp));
	if (!mp) {
		i_modsql_error("not enough memory");
		return -1;
	}
	conn->data = mp;
	return 0;
}

static void
s_mysql_destroy(struct vmod_sql_connection *conn)
{
	struct vmod_mysql_data *mp = conn->data;
	free(mp->buffer);
	free(mp->mysql);
	free(mp);
	conn->data = NULL;
}
  

static int
s_mysql_connect(struct vmod_sql_connection *conn)
{
	struct vmod_mysql_data *mp = conn->data;
	char *host, *socket_name = NULL;
	char *s;
	int port = 0;
	
	mp->mysql = malloc(sizeof(MYSQL));
	if (!mp->mysql) {
		i_modsql_error("not enough memory");
		conn->state = state_disabled;
		return -1;
	}
	
	mysql_init(mp->mysql);

	host = i_modsql_findparam(conn->param, "server");
	if (host && host[0] == '/') {
		host = "localhost";
		socket_name = host;
	}

	s = i_modsql_findparam(conn->param, "port");
	if (s)
		port = atoi(s);

	s = i_modsql_findparam(conn->param, "config");
	if (s)
		mysql_options(mp->mysql, MYSQL_READ_DEFAULT_FILE, s);

	s = i_modsql_findparam(conn->param, "group");
	if (s)
		mysql_options(mp->mysql, MYSQL_READ_DEFAULT_GROUP, s);

	s = i_modsql_findparam(conn->param, "cacert");
	if (s)
		mysql_ssl_set(mp->mysql, NULL, NULL, s, NULL, NULL);
	debug(conn, 1, ("connecting to database"));
	if (!mysql_real_connect(mp->mysql, 
				host,
				i_modsql_findparam(conn->param, "user"),
				i_modsql_findparam(conn->param, "password"),
				i_modsql_findparam(conn->param, "database"),
				port,
				socket_name,
				CLIENT_MULTI_RESULTS)) {
		i_modsql_error("cannot connect: %s", mysql_error(mp->mysql));
		return -1;
	}
	debug(conn, 1, ("connected to database"));
  
	return 0;
}

static int 
s_mysql_disconnect(struct vmod_sql_connection *conn)
{
	struct vmod_mysql_data *mp = conn->data;
	mysql_close(mp->mysql);
	return 0;
}

static int
s_mysql_query(struct vmod_sql_connection *conn, const char *query)
{
	struct vmod_mysql_data *mp = conn->data;
	int rc;
	int i;
	MYSQL *mysql;
	
	for (i = 0; i < 10; i++) {
		mysql = mp->mysql;
		rc = mysql_query(mysql, query);
		if (rc) {
			check_errno(conn);
			if (conn->state != state_init)
				return -1;
			i_modsql_connect(conn);
			if (conn->state != state_connected)
				return -1;
			continue;
		}
		mp->result = mysql_store_result(mp->mysql);
		if (mp->result) {
			conn->state = state_result;
			rc = 0;
		} else if (mysql_field_count(mp->mysql)) {
			i_modsql_error("cannot store result: %s",
				   mysql_error(mp->mysql));
			conn->state = state_error;
			rc = 1;
		} else
			rc = 0;
		break;
	}
	return rc;
}

static int
s_mysql_free_result(struct vmod_sql_connection *conn)
{
	struct vmod_mysql_data *mp = conn->data;
	mysql_free_result(mp->result);
	mp->result = NULL;
	return 0;
}

static unsigned
s_mysql_num_fields(struct vmod_sql_connection *conn)
{
	struct vmod_mysql_data *mp = conn->data;
	return mysql_num_fields(mp->result);
}

static unsigned
s_mysql_num_tuples(struct vmod_sql_connection *conn)
{
	struct vmod_mysql_data *mp = conn->data;
	return mysql_num_rows(mp->result);
}

static const char *
s_mysql_get_column(struct vmod_sql_connection *conn, size_t nrow, size_t ncol)
{
	struct vmod_mysql_data *mp = conn->data;
	MYSQL_ROW row;

	if (nrow >= mysql_num_rows(mp->result) ||
	    ncol >= mysql_num_fields (mp->result))
		return NULL;
  
	mysql_data_seek(mp->result, nrow);
	row = mysql_fetch_row(mp->result);
	if (!row)
		return NULL;
	return row[ncol];
}

static char *
s_mysql_escape(struct vmod_sql_connection *conn, const char *arg)
{
	struct vmod_mysql_data *mp = conn->data;
	size_t len, size;
	char *p;

	switch (conn->state) {
	case state_connected:
	case state_result:
		break;

	case state_error:
	case state_init:
		if (i_modsql_connect(conn))
			return NULL;
		break;

	default:
		return NULL;
	}

	len = strlen(arg);
	size = 2 * len + 1;
	if (mp->bufsize < size) {
		p = realloc(mp->buffer, size);
		if (!p) {
			i_modsql_error("not enough memory");
			return NULL;
		}
		mp->buffer = p;
		mp->bufsize = size;
	}
			 
	mysql_real_escape_string(mp->mysql, mp->buffer, arg, len);

	p = strdup(mp->buffer);
	if (!p)
		i_modsql_error("not enough memory");
	return p;
}

long
s_mysql_affected_rows(struct vmod_sql_connection *conn)
{
	struct vmod_mysql_data *mp = conn->data;
	return mysql_affected_rows(mp->mysql);
}
		
struct vmod_sql_backend i_modsql_mysql_backend = {
	"mysql",
	s_mysql_init,
	s_mysql_destroy,
	s_mysql_connect,
	s_mysql_disconnect,
	s_mysql_escape,
	s_mysql_query,
	s_mysql_num_tuples,
	s_mysql_num_fields,
	s_mysql_free_result,
	s_mysql_get_column,
	s_mysql_affected_rows
};