#include "mem.h"
#include "ssb.h"
#include "ssb.db.h"
#include "ssb.import.h"
#include "ssb.import.h"
#include "ssb.export.h"
#include "task.h"
#include "taskstub.js.h"
#include "tests.h"
#include "util.js.h"

#include <memcheck.h>
#include <quickjs-libc.h>
#include <quickjs.h>
#include <sqlite3.h>
#include <xopt.h>

#include <string.h>

#if !defined(_WIN32) && !defined(__MACH__)
#include <signal.h>
#include <sys/prctl.h>
#include <sys/resource.h>
#include <unistd.h>
#endif

#if !defined(_countof)
#define _countof(a) ((int)(sizeof((a)) / sizeof(*(a))))
#endif

#define XOPT_PARSE(name, flags, options, config_ptr, argc, argv, extrac_ptr, extrav_ptr, err_ptr, autohelp_file, autohelp_usage, autohelp_prefix, autohelp_suffix, autohelp_spacer) do { \
		xoptContext *_xopt_ctx; \
		*(err_ptr) = NULL; \
		_xopt_ctx = xopt_context((name), (options), ((flags) ^ XOPT_CTX_POSIXMEHARDER), (err_ptr)); \
		if (*(err_ptr)) break; \
		*extrac_ptr = xopt_parse(_xopt_ctx, (argc), (argv), (config_ptr), (extrav_ptr), (err_ptr)); \
		if ((config_ptr)->help) \
		{ \
			xoptAutohelpOptions __xopt_autohelp_opts; \
			__xopt_autohelp_opts.usage = (autohelp_usage); \
			__xopt_autohelp_opts.prefix = (autohelp_prefix); \
			__xopt_autohelp_opts.suffix = (autohelp_suffix); \
			__xopt_autohelp_opts.spacer = (autohelp_spacer); \
			xopt_autohelp(_xopt_ctx, (autohelp_file), &__xopt_autohelp_opts, (err_ptr)); \
			if (*(err_ptr)) goto __xopt_end_free_extrav; \
			free(_xopt_ctx); \
			goto xopt_help; \
		} \
		if (*(err_ptr)) goto __xopt_end_free_ctx; \
	__xopt_end_free_ctx: \
		free(_xopt_ctx); \
		break; \
	__xopt_end_free_extrav: \
		free(*(extrav_ptr)); \
		free(_xopt_ctx); \
		break; \
	} while (false)

static int _tf_command_test(const char* file, int argc, char* argv[]);
static int _tf_command_import(const char* file, int argc, char* argv[]);
static int _tf_command_export(const char* file, int argc, char* argv[]);
static int _tf_command_run(const char* file, int argc, char* argv[]);
static int _tf_command_sandbox(const char* file, int argc, char* argv[]);
static int _tf_command_post(const char* file, int argc, char* argv[]);
static int _tf_command_check(const char* file, int argc, char* argv[]);
static int _tf_command_usage(const char* file, int argc, char* argv[]);

typedef struct _command_t {
	const char* name;
	int (*callback)(const char* file, int argc, char* argv[]);
	const char* description;
} command_t;

const command_t k_commands[] = {
	{ "run", _tf_command_run, "Run tildefriends (default)." },
	{ "sandbox", _tf_command_sandbox, "Run a sandboxed tildefriends sandbox process (used internally)." },
	{ "post", _tf_command_post, "Create an SSB post." },
	{ "import", _tf_command_import, "Import apps to SSB." },
	{ "export", _tf_command_export, "Export apps from SSB." },
	{ "test", _tf_command_test, "Test SSB." },
	{ "check", _tf_command_check, "Validate messages in the SSB database." },
};

void shedPrivileges()
{
#if !defined(_WIN32)
	struct rlimit zeroLimit;
	zeroLimit.rlim_cur = 0;
	zeroLimit.rlim_max = 0;

	// RLIMIT_AS
	// RLIMIT_CORE
	// RLIMIT_CPU
	// RLIMIT_DATA
	// RLIMIT_FSIZE
	// RLIMIT_RSS
	// RLIMIT_RTPRIO
	// RLIMIT_RTTIME
	// RLIMIT_SIGPENDING
	// RLIMIT_STACK

	if (setrlimit(RLIMIT_FSIZE, &zeroLimit) != 0)
	{
		perror("setrlimit(RLIMIT_FSIZE, {0, 0})");
		exit(-1);
	}
	if (setrlimit(RLIMIT_NOFILE, &zeroLimit) != 0)
	{
		perror("setrlimit(RLIMIT_NOFILE, {0, 0})");
		exit(-1);
	}
	if (setrlimit(RLIMIT_NPROC, &zeroLimit) != 0)
	{
		perror("setrlimit(RLIMIT_NPROC, {0, 0})");
		exit(-1);
	}
#if !defined(__MACH__)
	if (setrlimit(RLIMIT_LOCKS, &zeroLimit) != 0)
	{
		perror("setrlimit(RLIMIT_LOCKS, {0, 0})");
		exit(-1);
	}
	if (setrlimit(RLIMIT_MSGQUEUE, &zeroLimit) != 0)
	{
		perror("setrlimit(RLIMIT_MSGQUEUE, {0, 0})");
		exit(-1);
	}
#endif
#endif
}

static int _tf_command_test(const char* file, int argc, char* argv[])
{
	typedef struct args_t {
		const char* tests;
		bool help;
	} args_t;

	xoptOption options[] = {
		{ "tests", 't', offsetof(args_t, tests), NULL, XOPT_TYPE_STRING, NULL, "Comma-separated list of test names to run." },
		{ "help", 'h', offsetof(args_t, help), NULL, XOPT_TYPE_BOOL, NULL, "Shows this help message." },
		XOPT_NULLOPTION,
	};

	args_t args = { 0 };
	const char** extras = NULL;
	int extra_count = 0;
	const char *err = NULL;
	XOPT_PARSE(file, XOPT_CTX_KEEPFIRST | XOPT_CTX_STRICT, options, &args, argc, (const char**)argv, &extra_count, &extras, &err, stderr, "test [options]", "options:", NULL, 15);
	if (extras)
	{
		free((void*)extras);
	}
	if (err)
	{
		fprintf(stderr, "Error: %s\n", err);
		return 2;
	}

	tf_test_options_t test_options =
	{
		.exe_path = file,
		.tests = args.tests,
	};
	tf_tests(&test_options);
	return 0;
xopt_help:
	if (extras)
	{
		free((void*)extras);
	}
	return 1;
}

static int _tf_command_import(const char* file, int argc, char* argv[])
{
	typedef struct args_t {
		const char* user;
		const char* db_path;
		bool help;
	} args_t;

	xoptOption options[] = {
		{ "user", 'u', offsetof(args_t, user), NULL, XOPT_TYPE_STRING, NULL, "User into whose account apps will be imported (default: \"import\")." },
		{ "db-path", 'd', offsetof(args_t, db_path), NULL, XOPT_TYPE_STRING, NULL, "Sqlite database path (default: db.sqlite)." },
		{ "help", 'h', offsetof(args_t, help), NULL, XOPT_TYPE_BOOL, NULL, "Shows this help message." },
		XOPT_NULLOPTION,
	};

	args_t args = { .user = "import" };
	const char** extras = NULL;
	int extra_count = 0;
	const char *err = NULL;
	XOPT_PARSE(file, XOPT_CTX_KEEPFIRST | XOPT_CTX_POSIXMEHARDER | XOPT_CTX_STRICT, options, &args, argc, (const char**)argv, &extra_count, &extras, &err, stderr, "import [options] [paths] ...", "options:", NULL, 15);

	if (err)
	{
		fprintf(stderr, "Error: %s\n", err);
		if (extras)
		{
			free((void*)extras);
		}
		return 2;
	}

	sqlite3* db = NULL;
	if (args.db_path)
	{
		sqlite3_open(args.db_path, &db);
	}
	tf_ssb_t* ssb = tf_ssb_create(NULL, NULL, db, NULL);
	if (extra_count)
	{
		for (int i = 0; i < extra_count; i++)
		{
			printf("Importing %s...\n", extras[i]);
			tf_ssb_import(ssb, args.user, extras[i]);
		}
	}
	else
	{
		printf("Importing %s...\n", "apps");
		tf_ssb_import(ssb, args.user, "apps");
	}
	tf_ssb_destroy(ssb);
	if (db)
	{
		sqlite3_close(db);
	}

	if (extras)
	{
		free((void*)extras);
	}
	return 0;

xopt_help:
	if (extras)
	{
		free((void*)extras);
	}
	return 1;
}

static int _tf_command_export(const char* file, int argc, char* argv[])
{
	typedef struct args_t {
		const char* db_path;
		bool help;
	} args_t;

	xoptOption options[] = {
		{ "db-path", 'd', offsetof(args_t, db_path), NULL, XOPT_TYPE_STRING, NULL, "Sqlite database path (default: db.sqlite)." },
		{ "help", 'h', offsetof(args_t, help), NULL, XOPT_TYPE_BOOL, NULL, "Shows this help message." },
		XOPT_NULLOPTION,
	};

	args_t args = { 0 };
	const char** extras = NULL;
	int extra_count = 0;
	const char *err = NULL;
	XOPT_PARSE(file, XOPT_CTX_KEEPFIRST | XOPT_CTX_POSIXMEHARDER | XOPT_CTX_STRICT, options, &args, argc, (const char**)argv, &extra_count, &extras, &err, stderr, "export [options] [paths] ...", "options:", NULL, 15);

	if (err)
	{
		fprintf(stderr, "Error: %s\n", err);
		if (extras)
		{
			free((void*)extras);
		}
		return 2;
	}
	sqlite3* db = NULL;
	if (args.db_path)
	{
		sqlite3_open(args.db_path, &db);
	}
	tf_ssb_t* ssb = tf_ssb_create(NULL, NULL, db, NULL);
	if (extra_count)
	{
		for (int i = 0; i < extra_count; i++)
		{
			printf("Exporting %s...\n", extras[i]);
			tf_ssb_export(ssb, extras[i]);
		}
	}
	else
	{
		const char* k_export[] = {
			"/~cory/api",
			"/~cory/admin",
			"/~cory/apps",
			"/~cory/db",
			"/~cory/docs",
			"/~cory/follow",
			"/~cory/ssb",
		};
		for (int i = 0; i < (int)_countof(k_export); i++)
		{
			printf("Exporting %s...\n", k_export[i]);
			tf_ssb_export(ssb, k_export[i]);
		}
	}
	tf_ssb_destroy(ssb);
	if (db)
	{
		sqlite3_close(db);
	}

	if (extras)
	{
		free((void*)extras);
	}
	return 0;

xopt_help:
	if (extras)
	{
		free((void*)extras);
	}
	return 1;
}

typedef struct tf_run_args_t {
	const char* script;
	int ssb_port;
	int http_port;
	int https_port;
	const char* db_path;
	const char* secrets_path;
	int count;
	const char* args;
	bool help;
} tf_run_args_t;

typedef struct _tf_run_thread_data_t
{
	tf_run_args_t args;
	int index;
	int result;
} tf_run_thread_data_t;

static int _tf_run_task(const tf_run_args_t* args, int index)
{
	int result = -1;
	tf_task_t* task = tf_task_create();
	tf_task_set_trusted(task, true);
	tf_task_set_ssb_port(task, args->ssb_port ? args->ssb_port + index : 0);
	tf_task_set_http_port(task, args->http_port ? args->http_port + index : 0);
	tf_task_set_https_port(task, args->https_port ? args->https_port + index : 0);
	tf_task_set_args(task, args->args);
	const char* db_path = args->db_path;
	const char* secrets_path = args->secrets_path;
	char db_path_buffer[256];
	char secrets_path_buffer[256];
	if (index)
	{
		snprintf(db_path_buffer, sizeof(db_path_buffer), "%s.%d", args->db_path, index);
		db_path = db_path_buffer;
		snprintf(secrets_path_buffer, sizeof(secrets_path_buffer), "%s.%d", args->secrets_path, index);
		secrets_path = secrets_path_buffer;
	}
	tf_task_set_db_path(task, db_path);
	tf_task_set_secrets_path(task, secrets_path);
	tf_task_activate(task);
	if (args->ssb_port)
	{
		tf_ssb_import(tf_task_get_ssb(task), "core", "apps");
	}
	if (tf_task_execute(task, args->script))
	{
		tf_task_run(task);
		result = 0;
	}
	tf_task_destroy(task);
	return result;
}

static void _tf_run_task_thread(void* data)
{
	tf_run_thread_data_t* info = data;
	info->result = _tf_run_task(&info->args, info->index);
}

static int _tf_command_run(const char* file, int argc, char* argv[])
{
	xoptOption options[] = {
		{ "script", 's', offsetof(tf_run_args_t, script), NULL, XOPT_TYPE_STRING, NULL, "Script to run (default: core/core.js)." },
		{ "ssb-port", 'b', offsetof(tf_run_args_t, ssb_port), NULL, XOPT_TYPE_INT, NULL, "Port on which to run SSB (default: 8009)." },
		{ "http-port", 'p', offsetof(tf_run_args_t, http_port), NULL, XOPT_TYPE_INT, NULL, "Port on which to run Tilde Friends web server (default: 12345)." },
		{ "https-port", 'q', offsetof(tf_run_args_t, https_port), NULL, XOPT_TYPE_INT, NULL, "Port on which to run secure Tilde Friends web server (default: 12346)." },
		{ "db-path", 'd', offsetof(tf_run_args_t, db_path), NULL, XOPT_TYPE_STRING, NULL, "Sqlite database path (default: db.sqlite)." },
		{ "secrets-path", 'i', offsetof(tf_run_args_t, secrets_path), NULL, XOPT_TYPE_STRING, NULL, "Secrets/identity path." },
		{ "count", 'n', offsetof(tf_run_args_t, count), NULL, XOPT_TYPE_INT, NULL, "Number of instances to run." },
		{ "args", 'a', offsetof(tf_run_args_t, args), NULL, XOPT_TYPE_STRING, NULL, "Arguments of the form key=value,foo=bar,verbose=true." },
		{ "help", 'h', offsetof(tf_run_args_t, help), NULL, XOPT_TYPE_BOOL, NULL, "Shows this help message." },
		XOPT_NULLOPTION,
	};

	tf_run_args_t args =
	{
		.count = 1,
		.script = "core/core.js",
		.http_port = 12345,
		.https_port = 12346,
		.ssb_port = 8009,
		.db_path = "db.sqlite",
		.secrets_path = "/.config/tildefriends/secret",
	};
	const char** extras = NULL;
	int extra_count = 0;
	const char *err = NULL;
	XOPT_PARSE(file, XOPT_CTX_KEEPFIRST | XOPT_CTX_POSIXMEHARDER | XOPT_CTX_STRICT, options, &args, argc, (const char**)argv, &extra_count, &extras, &err, stderr, "run [options] [paths] ...", "options:", NULL, 15);

	if (err)
	{
		fprintf(stderr, "Error: %s\n", err);
		if (extras)
		{
			free((void*)extras);
		}
		return 2;
	}
	if (extras)
	{
		free((void*)extras);
	}

	int result = 0;
#if !defined(_WIN32) && !defined(__MACH__)
	setpgid(0, 0);
#endif

	if (args.count == 1)
	{
		_tf_run_task(&args, 0);
	}
	if (args.count > 1)
	{
		uv_thread_t* threads = tf_malloc(sizeof(uv_thread_t) * args.count);
		tf_run_thread_data_t* data = tf_malloc(sizeof(tf_run_thread_data_t) * args.count);
		for (int i = 0 ; i < args.count; i++)
		{
			data[i] = (tf_run_thread_data_t)
			{
				.args = args,
				.index = i,
			};
			uv_thread_create(&threads[i], _tf_run_task_thread, &data[i]);
		}
		for (int i = 0; i < args.count; i++)
		{
			uv_thread_join(&threads[i]);
			if (data[i].result != 0)
			{
				result = data[i].result;
			}
		}
		tf_free(data);
		tf_free(threads);
	}

	return result;

xopt_help:
	if (extras)
	{
		free((void*)extras);
	}
	return 1;
}

static int _tf_command_sandbox(const char* file, int argc, char* argv[])
{
	typedef struct args_t {
		const char* script;
		bool help;
	} args_t;

	xoptOption options[] = {
		{ "help", 'h', offsetof(args_t, help), NULL, XOPT_TYPE_BOOL, NULL, "Shows this help message." },
		XOPT_NULLOPTION,
	};

	args_t args = { 0 };
	const char** extras = NULL;
	int extra_count = 0;
	const char *err = NULL;
	XOPT_PARSE(file, XOPT_CTX_KEEPFIRST | XOPT_CTX_POSIXMEHARDER | XOPT_CTX_STRICT, options, &args, argc, (const char**)argv, &extra_count, &extras, &err, stderr, "sandbox [options]", "options:", NULL, 15);

	if (err)
	{
		fprintf(stderr, "Error: %s\n", err);
		if (extras)
		{
			free((void*)extras);
		}
		return 2;
	}
	if (extras)
	{
		free((void*)extras);
	}

#if !defined(_WIN32) && !defined(__MACH__)
	prctl(PR_SET_PDEATHSIG, SIGHUP);
#endif
	tf_task_t* task = tf_task_create();
	tf_task_configure_from_stdin(task);
	shedPrivileges();
	/* The caller will trigger tf_task_activate with a message. */
	tf_task_run(task);
	tf_task_destroy(task);
	return 0;

xopt_help:
	if (extras)
	{
		free((void*)extras);
	}
	return 1;
}

static int _tf_command_post(const char* file, int argc, char* argv[])
{
	typedef struct args_t {
		char* message;
		bool help;
	} args_t;

	xoptOption options[] = {
		{ "message", 'm', offsetof(args_t, message), NULL, XOPT_REQUIRED | XOPT_TYPE_STRING, "TEXT", "Text to post." },
		{ "help", 'h', offsetof(args_t, help), NULL, XOPT_TYPE_BOOL, NULL, "Shows this help message." },
		XOPT_NULLOPTION,
	};

	args_t args = { 0 };
	const char** extras = NULL;
	int extra_count = 0;
	const char *err = NULL;
	XOPT_PARSE(file, XOPT_CTX_KEEPFIRST | XOPT_CTX_STRICT, options, &args, argc, (const char**)argv, &extra_count, &extras, &err, stderr, "post [options]", "options:", NULL, 15);
	if (extras)
	{
		free((void*)extras);
	}
	if (err)
	{
		fprintf(stderr, "Error: %s\n", err);
		return 2;
	}

	tf_ssb_t* ssb = tf_ssb_create(NULL, NULL, NULL, NULL);
	tf_ssb_broadcast_listener_start(ssb, false);
	tf_ssb_append_post(ssb, args.message);
	tf_ssb_destroy(ssb);
	return 0;

xopt_help:
	if (extras)
	{
		free((void*)extras);
	}
	return 1;
}

static int _tf_command_check(const char* file, int argc, char* argv[])
{
	typedef struct args_t {
		bool help;
	} args_t;

	xoptOption options[] = {
		{ "help", 'h', offsetof(args_t, help), NULL, XOPT_TYPE_BOOL, NULL, "Shows this help message." },
		XOPT_NULLOPTION,
	};

	args_t args = { 0 };
	const char** extras = NULL;
	int extra_count = 0;
	const char *err = NULL;
	XOPT_PARSE(file, XOPT_CTX_KEEPFIRST | XOPT_CTX_STRICT, options, &args, argc, (const char**)argv, &extra_count, &extras, &err, stderr, "post [options]", "options:", NULL, 15);
	if (err)
	{
		if (extras)
		{
			free((void*)extras);
		}
		fprintf(stderr, "Error: %s\n", err);
		return 2;
	}

	bool result = true;
	sqlite3* db = NULL;
	sqlite3_open("db.sqlite", &db);
	if (extra_count)
	{
		for (int i = 0; i < extra_count; i++)
		{
			result = result && tf_ssb_db_check(db, extras[i]);
		}
	}
	else
	{
		result = tf_ssb_db_check(db, NULL);
	}
	sqlite3_close(db);
	if (extras)
	{
		free((void*)extras);
	}
	return result ? EXIT_SUCCESS : EXIT_FAILURE;

xopt_help:
	if (extras)
	{
		free((void*)extras);
	}
	return 1;
}

static int _tf_command_usage(const char* file, int argc, char* argv[])
{
	printf("Usage: %s command [command-options]\n", file);
	printf("commands:\n");
	for (int i = 0; i < (int)_countof(k_commands); i++)
	{
		printf("  %s - %s\n", k_commands[i].name, k_commands[i].description);
	}
	return 0;
}

#if !defined(_WIN32)
static void _do_leak_checks(int sig)
{
	VALGRIND_DO_LEAK_CHECK;
}
#endif

int main(int argc, char* argv[])
{
#if !defined(_WIN32)
	prctl(PR_SET_PDEATHSIG, SIGKILL);
#endif
	tf_mem_replace_uv_allocator();
	tf_mem_replace_tls_allocator();
	uv_setup_args(argc, argv);
	tf_taskstub_startup();

#if !defined(_WIN32)
	if (signal(SIGPIPE, SIG_IGN) == SIG_ERR)
	{
		perror("signal");
	}
	if (signal(SIGHUP, _do_leak_checks) == SIG_ERR)
	{
		perror("signal");
	}
#endif

	if (argc >= 2)
	{
		for (int i = 0; i < (int)_countof(k_commands); i++)
		{
			const command_t* command = &k_commands[i];
			if (strcmp(argv[1], command->name) == 0)
			{
				return command->callback(argv[0], argc - 2, argv + 2);
			}
		}
		return _tf_command_usage(argv[0], argc, argv);
	}

	return _tf_command_run(argv[0], argc - 1, argv + 1);
}