/*-------------------------------------------------------------------------
 *
 * pg_snapclone_master.c
 *
 * Portions Copyright (c) 1996-2008, PostgreSQL Global Development Group
 * Portions Copyright (c) 1994, Regents of the University of California
 *
 * $PostgreSQL$
 *
 * Original author:		Simon Riggs  simon@2ndquadrant.com
 * Current maintainer:	Simon Riggs
 *-------------------------------------------------------------------------
 */
#include "postgres_fe.h"
#include "libpq-fe.h"
#include "miscadmin.h"

#include <fcntl.h>
#include <signal.h>

#ifdef WIN32
int			getopt(int argc, char *const argv[], const char *optstring);
#else
#include <unistd.h>

#ifdef HAVE_GETOPT_H
#include <getopt.h>
#endif
#endif   /* ! WIN32 */

extern char *optarg;
extern int	optind;

/* Options and defaults */
int			sleeptime = 1;
int			waittime = -1;		/* how long we have been waiting, -1 no wait
								 * yet */
int			maxwaittime = 0;	/* how long are we prepared to wait for? */
bool		triggered = false;	/* have we been triggered? */

#define 	MAX_STATEMENT_SIZE			8192

#define		DEFAULT_PORT		"5432"
#define		DEFAULT_USER		"postgres"
#define		DEFAULT_DATABASE	"postgres"

#define		DEFAULT_SNAPSHOT_FILEPATH	"pg_snapclone.pid"

/* Database */
char 		query[MAX_STATEMENT_SIZE];

struct	passwd *password_file_entry;

PGconn	   *dbconn;
PGresult   *dbresult;

int			fd;

char	   *snapshot_filepath = DEFAULT_SNAPSHOT_FILEPATH;
int			pid;


static void help(const char *progname);
static void sighandler(int sig);


static void
sighandler(int sig)
{
	triggered = true;
}

static bool
SnapshotFilePresent(void)
{
	fd = open(snapshot_filepath, O_RDONLY);

	if (fd != -1)
	{
		close(fd);
		return true;
	}

	return false;
}

int
main(int argc, char **argv)
{
	int			c;

	static struct option long_opts[] = {
		{"host", required_argument, NULL, 'h'},
		{"port", required_argument, NULL, 'p'},
		{"username", required_argument, NULL, 'U'},
		{"password", no_argument, NULL, 'W'},
		{"dbname", required_argument, NULL, 'd'},
		{"help", no_argument, NULL, 'H'},
		{NULL, 0, NULL, 0}
	};

	const char *progname;

	const char *dbname = NULL;
	char	   *host = NULL;
	char	   *port = NULL;
	char	   *username = NULL;
	bool		password = false;
	bool		usage_error = false;

	pid = getpid();
	progname = get_progname(argv[0]);

	while ((c = getopt_long(argc, argv, "h:Hp:U:Wd:f:w:u", long_opts, NULL)) != -1)
	{
		switch (c)
		{
			case 'h':
				host = optarg;
				break;
			case 'p':
				port = optarg;
				break;
			case 'U':
				username = optarg;
				break;
			case 'W':
				password = true;
				break;
			case 'd':
				dbname = optarg;
				break;
			case 'f':
				snapshot_filepath = optarg;
				break;
			case 'w':			/* Max wait time */
				maxwaittime = atoi(optarg);
				if (maxwaittime < 0)
				{
					fprintf(stderr, "usage: -w maxwaittime incorrectly set\n");
					help(progname);
					exit(2);
				}
				break;
			case 'H':
			default:
				usage_error = true;
				break;
		}
	}

	(void) signal(SIGINT, sighandler);
	(void) signal(SIGQUIT, sighandler);
	(void) signal(SIGTERM, sighandler);

	/* additional argument overrides -d option (if present) to set database */
	if ((argc - optind) == 1)
		dbname = argv[optind];

	if (usage_error)
	{
		help(progname);
		exit(2);
	}	

	/*
	 * Is SnapshotFile already there?
	 */
	fd = open(snapshot_filepath, O_RDONLY);		/* we hope this fails */

	if (fd != -1)
	{
		/* read file to get pid */
		if (read(fd, &pid, sizeof(pid)) != sizeof(pid))
			fprintf(stderr, "\nsnapshot file \"%s\" exists, but has read errors\n", 
						snapshot_filepath);
		else
			fprintf(stderr, "\nsnapshot file \"%s\" exists already for pid=%d\n",
						snapshot_filepath, pid);
		close(fd);
		exit(2);
	}

	if (dbname == NULL)
	{
		if (getenv("PGDATABASE"))
			dbname = getenv("PGDATABASE");
		else if (getenv("PGUSER"))
			dbname = getenv("PGUSER");
		else
			dbname = DEFAULT_DATABASE;
	}

	if (username == NULL)
	{
		if (getenv("PGUSER"))
			username = getenv("PGUSER");
		else
		{
			/* try to connect to a database same as current username */
	        password_file_entry = getpwuid(geteuid());
	        if (password_file_entry)
				username = password_file_entry->pw_name;
			else
				username = DEFAULT_USER;
		}
	}

	if (port == NULL)
	{
		if (getenv("PGPORT"))
			port = getenv("PGPORT");
		else
            port = DEFAULT_PORT;
	}

	/*
	 * Connect
	 */
    snprintf(query, MAX_STATEMENT_SIZE,
                         "dbname=%s port=%s user=%s", 
							dbname, port, username);
    dbconn = PQconnectdb(query);
	if (PQstatus(dbconn) != CONNECTION_OK)
	{
		fprintf(stderr, "\n%s:  failed to connect: %s\n", progname, query);
		exit(2);
	}

	/*
	 * BEGIN transaction
	 */
	dbresult = PQexec(dbconn, 
						"BEGIN ISOLATION LEVEL SERIALIZABLE, READ ONLY");
    if (PQresultStatus(dbresult) != PGRES_COMMAND_OK)
	{
		fprintf(stderr, "\n%s:  failed to start serializable transaction\n", progname);
		PQclear(dbresult);
	    PQfinish(dbconn);
		exit(2);
	}
	PQclear(dbresult);

	/*
	 * Publish snapshot
	 */
    snprintf(query, MAX_STATEMENT_SIZE,
                         "select pg_snapclone_publish()");
	dbresult = PQexec(dbconn, query);
    if (PQresultStatus(dbresult) != PGRES_TUPLES_OK)
	{
		fprintf(stderr, "\n%s:  failed to publish snapshot", progname);
		fprintf(stderr, "\n%s:  have you installed pg_snapclone_tools?\n", progname);
		PQclear(dbresult);
	    PQfinish(dbconn);
		exit(2);
	}
	pid = atoi(PQgetvalue(dbresult,0,0));
	PQclear(dbresult);

	/*
	 * Write SnapshotFile
	 */
	fd = open(snapshot_filepath, O_RDWR | O_CREAT, S_IRWXU | S_IRGRP);

	if (fd < 0)
	{
		/* report snapshot_filepath already exists */
		fprintf(stderr, "\n%s:  unable to open snapshot file for write (errno=%d)", 
					progname, errno);
   	    exit(2);
	}

	if (write(fd, &pid, sizeof(pid)) != sizeof(pid))
	{
		/* report snapshot_filepath already exists */
		fprintf(stderr, "\n%s:  write error on snapshot file %s", 
					progname, snapshot_filepath);
		close(fd);
   	    exit(2);
	}
	close(fd);

	/*
	 * Main wait loop
	 */
	while (SnapshotFilePresent() && !triggered)
	{
		pg_usleep(sleeptime * 1000000L);

		waittime += sleeptime;

		if (waittime >= maxwaittime && maxwaittime > 0)
		{
			triggered = true;
			fprintf(stderr, "\n%s:  timed out after %d seconds\n", 
						progname, waittime);
			unlink(snapshot_filepath);
		}
	}

	/*
	 * COMMIT
	 */
	dbresult = PQexec(dbconn, "ROLLBACK");
    if (PQresultStatus(dbresult) != PGRES_COMMAND_OK)
	{
		fprintf(stderr, "\n%s:  failed to rollback\n", progname);
		PQclear(dbresult);
	    PQfinish(dbconn);
		exit(2);
	}
	PQclear(dbresult);

	/*
	 * Unpublish snapshot
	 */
    snprintf(query, MAX_STATEMENT_SIZE,
                         "select pg_snapclone_unpublish()");
	dbresult = PQexec(dbconn, query);
    if (PQresultStatus(dbresult) != PGRES_TUPLES_OK)
	{
		fprintf(stderr, "\n%s:  failed to unpublish snapshot\n", progname);
		PQclear(dbresult);
	    PQfinish(dbconn);
		exit(2);
	}
	PQclear(dbresult);

	/*
	 * Disconnect
	 */
    PQfinish(dbconn);

	if (triggered)
		exit(1);
	
	return 0;
}

static void
help(const char *progname)
{
	fprintf(stderr, "\n%s provides a command line interface to clone Postgres snapshots\n", progname);
	fprintf(stderr, "Usage:\n");
	fprintf(stderr, "  %s [OPTION]... DBNAME\n",progname);
	fprintf(stderr, "\nOptions:\n");
	fprintf(stderr, "  -f SNAPCLONEFILE	defines file for sharing snapshot pid ("
						DEFAULT_SNAPSHOT_FILEPATH ")\n");
	fprintf(stderr, "  -w MAXWAITTIME	max seconds to wait for a file deletion (0=no limit)(default=0)\n");
	printf(_("\nConnection options:\n"));
	printf(_("  -h, --host=HOSTNAME          database server host or socket directory\n"));
	printf(_("  -p, --port=PORT              database server port\n"));
	printf(_("  -d, --dbname=DBNAME             database to vacuum\n"));
	printf(_("  -U, --username=USERNAME      user name to connect as\n"));
	printf(_("  -W, --password               force password prompt\n"));
	fflush(stderr);
}
