#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <unistd.h>
#include <string.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <stdbool.h>
#include <getopt.h>
#include <sys/signalfd.h>
#include <sys/eventfd.h>
#include <signal.h>
#include <systemd/sd-daemon.h>
#include <cap-ng.h>
#include <pwd.h>
#include <grp.h>

#include "main.h"
#include "uring.h"
#include "config.h"
#include "server.h"
#include "cfgdir.h"
#include "announce.h"
#include "systemd.h"
#include "igmp.h"
#include "idle.h"

#define DEFAULT_HOMEDIR_PATH "/home/david/intest"
#define DEFAULT_MAIN_CONFIG_FILE_PATH "./mcproxy.conf"

/* Global, not in struct cfg since they're too low-level */
bool exiting = false;
unsigned debug_mask = DBG_ERROR | DBG_INFO;

/* Local */
static bool daemonize = false;
static FILE *log_file = NULL;
static const char *log_file_path = NULL;
static struct cfg *cfghack = NULL;

#define ANSI_RED	"\x1B[0;31m"
#define ANSI_GREEN	"\x1B[0;32m"
#define ANSI_YELLOW	"\x1B[0;33m"
#define ANSI_BLUE	"\x1B[0;34m"
#define ANSI_MAGENTA	"\x1B[0;35m"
#define ANSI_GREY	"\x1B[0;38;5;245m"
#define ANSI_NORMAL	"\x1B[0m"

static void
msg(enum debug_lvl lvl, const char *fmt, va_list ap)
{
	static bool first = true;
	static bool use_colors = false;
	static bool sd_daemon = false;
	const char *color;
	const char *sd_lvl;

	assert_return(lvl != 0 && !empty_str(fmt) && ap);

	while (first) {
		int fd;
		const char *e;

		first = false;

		/* assume we're not launched by systemd when daemonized */
		if (daemonize) {
			sd_daemon = false;
			use_colors = false;
			break;
		}

		if (log_file) {
			sd_daemon = false;
			use_colors = false;
			break;
		}

		if (getenv("NO_COLOR")) {
			sd_daemon = false;
			use_colors = false;
			break;
		}

		fd = fileno(stderr);
		if (fd < 0) {
			/* Umm... */
			sd_daemon = true;
			use_colors = false;
			break;
		}

		if (!isatty(fd)) {
			sd_daemon = true;
			use_colors = false;
			break;
		}

		/* systemd wouldn't normally set TERM */
		e = getenv("TERM");
		if (!e) {
			sd_daemon = true;
			use_colors = false;
			break;
		}

		if (streq(e, "dumb")) {
			sd_daemon = false;
			use_colors = false;
			break;
		}

		sd_daemon = false;
		use_colors = true;
	}

	switch (lvl) {
	case DBG_ERROR:
		sd_lvl = SD_ERR;
		color = use_colors ? ANSI_RED : NULL;
		break;
	case DBG_VERBOSE:
		sd_lvl = SD_INFO;
		color = NULL;
		break;
	case DBG_INFO:
		sd_lvl = SD_NOTICE;
		color = NULL;
		break;
	default:
		sd_lvl = SD_DEBUG;
		color = use_colors ? ANSI_GREY : NULL;
		break;
	}

	if (sd_daemon)
		fprintf(stderr, sd_lvl);
	else if (color)
		fprintf(stderr, color);

	vfprintf(log_file ? log_file : stderr, fmt, ap);

	if (color)
		fprintf(stderr, ANSI_NORMAL);
}

void
__debug(enum debug_lvl lvl, const char *fmt, ...)
{
	va_list ap;

	assert_return(lvl != 0 && !empty_str(fmt));

	va_start(ap, fmt);
	msg(lvl, fmt, ap);
	va_end(ap);
}

__attribute__((noreturn)) void
__die(const char *fmt, ...)
{
	va_list ap;

	if (!empty_str(fmt)) {
		va_start(ap, fmt);
		msg(DBG_ERROR, fmt, ap);
		va_end(ap);
	} else
		error("fmt not set");

	sd_notifyf(0, "STATUS=Error, shutting down");
	exit(EXIT_FAILURE);
};

static void
cfg_free(struct uring_task *task)
{
	struct cfg *cfg = container_of(task, struct cfg, task);

	assert_return(task);

	debug(DBG_SIG, "called");
	systemd_delete(cfg);
	xfree(cfg->igmp_iface);
	cfg->igmp_iface = NULL;
	xfree(cfg->cfg_path);
	cfg->cfg_path = NULL;
	exiting = true;
	/* The cfg struct is free:d in main() */
}

enum mcfg_keys {
	MCFG_KEY_INVALID = 0,
	MCFG_KEY_IGMP,
	MCFG_KEY_IGMP_IFACE
};

struct cfg_key_value_map mcfg_key_map[] = {
	{
		.key_name = "igmp",
		.key_value = MCFG_KEY_IGMP,
		.value_type = CFG_VAL_TYPE_BOOL,
	}, {
		.key_name = "igmp_iface",
		.key_value = MCFG_KEY_IGMP_IFACE,
		.value_type = CFG_VAL_TYPE_STRING,
	}, {
		.key_name = NULL,
		.key_value = MCFG_KEY_INVALID,
		.value_type = CFG_VAL_TYPE_INVALID,
	}
};

static void
cfg_read(struct cfg *cfg)
{
	FILE *cfgfile;
	const char *path;
	char buf[4096];
	char *pos = buf;
	size_t rd = 0;
	size_t r;

	assert_return(cfg);

	if (cfg->cfg_path)
		path = cfg->cfg_path;
	else
		path = DEFAULT_MAIN_CONFIG_FILE_PATH;

	cfgfile = fopen(path, "re");
	if (!cfgfile) {
		/* ENOENT is only an error with an explicitly set path */
		if (errno == ENOENT && !cfg->cfg_path)
			return;
		else if (errno == ENOENT)
			die("main config file (%s) missing", path);
		else
			die("fopen(%s): %m", path);
	}

	debug(DBG_CFG, "opened main config file (%s)", path);

	while (rd < sizeof(buf)) {
		r = fread(pos, 1, sizeof(buf) - rd - 1, cfgfile);
		if (r == 0)
			break;
		rd += r;
		pos += r;
	}

	if (rd == 0)
		die("main config file (%s) zero size", path);

	if (rd >= sizeof(buf))
		die("main config file (%s) too large", path);

	fclose(cfgfile);
	*pos = '\0';
	pos = buf;

	if (!config_parse_header(cfg, path, "mcproxy", &pos))
		die("main config file (%s) invalid", path);

	while (true) {
		int key;
		const char *keyname;
		struct cfg_value value;

		if (!config_parse_line(cfg, path, &pos, mcfg_key_map,
				       &key, &keyname, &value))
			break;

		if (key == MCFG_KEY_INVALID)
			die("main config file (%s) invalid", path);

		debug(DBG_CFG, "main cfg: key %s", keyname);

		switch (key) {

		case MCFG_KEY_IGMP:
			cfg->do_igmp = value.boolean;
			break;

		case MCFG_KEY_IGMP_IFACE:
			cfg->igmp_iface = xstrdup(value.str);
			if (!cfg->igmp_iface)
				die("xstrdup: %m");

			break;

		case MCFG_KEY_INVALID:
		default:
			die("main config file (%s) invalid", path);
		}
	}
}

const struct {
	const char *name;
	unsigned val;
} debug_category_str[] = {
	{
		.name = "config",
		.val = DBG_CFG,
	},{
		.name = "refcount",
		.val = DBG_REF,
	},{
		.name = "malloc",
		.val = DBG_MALLOC,
	},{
		.name = "announce",
		.val = DBG_ANN,
	},{
		.name = "signal",
		.val = DBG_SIG,
	},{
		.name = "uring",
		.val = DBG_UR,
	},{
		.name = "server",
		.val = DBG_SRV,
	},{
		.name = "proxy",
		.val = DBG_PROXY,
	},{
		.name = "rcon",
		.val = DBG_RCON,
	},{
		.name = "idle",
		.val = DBG_IDLE,
	},{
		.name = "igmp",
		.val = DBG_IGMP,
	},{
		.name = "systemd",
		.val = DBG_SYSD,
	},{
		.name = "dns",
		.val = DBG_DNS,
	},{
		.name = NULL,
		.val = 0,
	}
};

__attribute__((noreturn)) static void
usage(int argc, char **argv, bool invalid)
{
	if (invalid)
		info("Invalid option(s)");

	info("Usage: %s [OPTIONS]\n"
	     "\n"
	     "Valid options:\n"
	     "  -H, --homedir=DIR\tlook for configuration files in DIR\n"
	     "  -u, --user=USER\trun as USER\n"
	     "  -D, --daemonize\trun in daemon mode (disables stderr output)\n"
	     "  -l, --logfile=FILE\tlog to FILE instead of stderr\n"
	     "  -h, --help\t\tprint this information\n"
	     "  -v, --verbose\t\tenable verbose logging\n"
	     "  -d, --debug=CATEGORY\tenable debugging for CATEGORY\n"
	     "\t\t\t(use \"list\" to see available categories,\n"
	     "\t\t\t or \"all\" to enable all categories)\n",
	     argv ? argv[0] : "mcproxy");

	exit(invalid ? EXIT_FAILURE : EXIT_SUCCESS);
}

static struct cfg *
cfg_init(int argc, char **argv)
{
	struct cfg *cfg;
	int c;
	unsigned i;

	assert_die(argc > 0 && argv, "invalid arguments");

	cfg = zmalloc(sizeof(*cfg));
	if (!cfg)
		die("malloc: %m");

	cfg->uid = geteuid();
	cfg->gid = getegid();
	uring_task_init(cfg, &cfg->task, "cfg", NULL, cfg_free);
	list_init(&cfg->servers);

	while (true) {
		int option_index = 0;
		static struct option long_options[] = {
			{ "homedir",	required_argument,	0, 'H' },
			{ "user",	required_argument,	0, 'u' },
			{ "daemonize",	no_argument,		0, 'D' },
			{ "logfile",	required_argument,	0, 'l' },
			{ "help",	no_argument,		0, 'h' },
			{ "verbose",	no_argument,		0, 'v' },
			{ "debug",	required_argument,	0, 'd' },
			{ 0,		0,			0,  0  }
		};

		c = getopt_long(argc, argv, ":H:u:Dl:hvd:",
				long_options, &option_index);
		if (c == -1)
			break;

		switch (c) {
		case 'H':
			cfg->homedir = optarg;
			break;

		case 'v':
			debug_mask |= DBG_VERBOSE;
			break;

		case 'D':
			daemonize = true;
			break;

		case 'l':
			log_file_path = optarg;
			break;

		case 'u': {
			struct passwd *pwd;

			errno = 0;
			pwd = getpwnam(optarg);
			if (!pwd) {
				if (errno == 0)
					errno = ESRCH;
				if (errno == ESRCH)
					die("failed to find user %s", optarg);
				else
					die("failed to find user %s (%m)", optarg);
			}

			debug(DBG_CFG, "asked to execute with uid %ji gid %ji",
			      (intmax_t)pwd->pw_uid,
			      (intmax_t)pwd->pw_gid);
			cfg->uid = pwd->pw_uid;
			cfg->gid = pwd->pw_gid;
			break;
		}

		case 'd':
			if (strcaseeq(optarg, "all")) {
				debug_mask = ~0;
				break;
			} else if (strcaseeq(optarg, "list")) {
				info("Debug categories:");
				info(" * all");
				for (i = 0; debug_category_str[i].name; i++)
					info(" * %s", debug_category_str[i].name);
				exit(EXIT_FAILURE);
			}

			for (i = 0; debug_category_str[i].name; i++) {
				if (strcaseeq(optarg, debug_category_str[i].name))
					break;
			}

			if (!debug_category_str[i].name)
				usage(argc, argv, true);

			debug_mask |= debug_category_str[i].val;
			break;

		case 'h':
			usage(argc, argv, false);

		default:
			usage(argc, argv, true);
		}

	}

	if (optind < argc)
		usage(argc, argv, true);

	if (!cfg->homedir)
		cfg->homedir = DEFAULT_HOMEDIR_PATH;

	return cfg;
}

static void
cfg_apply(struct cfg *cfg)
{
	if (!cfg)
		die("invalid arguments");

	if (cfg->uid == 0 || cfg->gid == 0)
		/* This catches both -u root and running as root without -u */
		die("Execution as root is not supported (use -u <someuser>)");

	capng_clear(CAPNG_SELECT_BOTH);
	if (capng_updatev(CAPNG_ADD,
			  CAPNG_EFFECTIVE | CAPNG_PERMITTED,
			  CAP_NET_RAW, CAP_NET_BIND_SERVICE, -1))
		die("capng_updatev failed");

	if (geteuid() != cfg->uid) {
		if (capng_change_id(cfg->uid,
				    cfg->gid,
				    CAPNG_DROP_SUPP_GRP | CAPNG_CLEAR_BOUNDING))
			die("capng_change_id failed");
	} else {
		if (capng_apply(CAPNG_SELECT_BOTH)) {
			capng_clear(CAPNG_SELECT_BOTH);
			if (capng_apply(CAPNG_SELECT_BOTH))
				die("capng_apply failed");
		}
		setgroups(0, NULL);
	}

	if (daemonize) {
		if (daemon(1, 0) < 0)
			die("daemon() failed: %m");
	}

	if (log_file_path) {
		log_file = fopen(log_file_path, "ae");
		if (!log_file)
			die("fopen(%s) failed: %m", log_file_path);
	}

	/*
	 * Do this after caps have been dropped to make sure we're not
	 * accessing a directory we should have permissions to.
	 */
	if (chdir(cfg->homedir))
		die("chdir(%s): %m", cfg->homedir);

	if (debug_enabled(DBG_VERBOSE)) {
		char *wd;

		wd = get_current_dir_name();
		verbose("Homedir: %s", wd ? wd : "<unknown>");
		free(wd);
	}
}

struct signalfd_ev {
	struct uring_task task;
	//struct signalfd_siginfo buf;
	struct cfg *cfg;
	uint64_t buf;
};

static void
signalfd_free(struct uring_task *task)
{
	struct signalfd_ev *sev = container_of(task, struct signalfd_ev, task);

	assert_return(task);

	debug(DBG_SIG, "called");
	sev->cfg->sev = NULL;
	xfree(sev);
}

static void
dump_tree(struct cfg *cfg)
{
	struct server *server;

	assert_return(cfg);

	if (!debug_enabled(DBG_REF))
		return;

	debug(DBG_REF, "\n\n");
	debug(DBG_REF, "Dumping Tree");
	debug(DBG_REF, "============");
	uring_task_refdump(&cfg->task);
	uring_refdump(cfg->uev);
	idle_refdump(cfg->idle);
	if (cfg->sev)
		uring_task_refdump(&cfg->sev->task);
	igmp_refdump(cfg->igmp);
	announce_refdump(cfg->aev);
	if (cfg->iev)
		cfgdir_refdump(cfg->iev);
	list_for_each_entry(server, &cfg->servers, list)
		server_refdump(server);
	debug(DBG_REF, "============");
	debug(DBG_REF, "\n\n");
}

static struct dns_async *hack_dns = NULL;

static void
signalfd_read(struct cfg *cfg, struct uring_task *task, int res)
{
	struct signalfd_ev *sev = container_of(task, struct signalfd_ev, task);
	struct server *server, *stmp;

	assert_return(cfg && task);
	assert_task_alive(DBG_SIG, task);

	if (res != sizeof(sev->buf))
		die("error in signalfd (%i)", res);

	if (sev->buf < 1000) {
		verbose("got a signal to quit");
		sd_notifyf(0, "STOPPING=1\nSTATUS=Received signal, exiting");
		exit(EXIT_SUCCESS);
	} else if (sev->buf < 10000) {
		verbose("got a signal to dump tree");
		sd_notifyf(0, "STOPPING=1\nSTATUS=Received signal, exiting");
		dump_tree(cfg);
		uring_task_put(cfg, &sev->task);
		igmp_delete(cfg);
		announce_delete(cfg);
		idle_delete(cfg);
		cfgdir_delete(cfg);
		list_for_each_entry_safe(server, stmp, &cfg->servers, list)
			server_delete(cfg, server);
		uring_delete(cfg);
		return;
	} else {
		debug(DBG_DNS, "DNS lookup complete, dns: %p, dns->cb: %p",
		      hack_dns,
		      hack_dns ? hack_dns->cb : NULL);

		if (!hack_dns || !hack_dns->cb) {
			error("DNS callback not set");
			goto out;
		}

		hack_dns->cb(hack_dns);
	}

out:
	uring_read(cfg, &sev->task, &sev->buf, sizeof(sev->buf), signalfd_read);
}

static int hack_efd = -1;

static void
hack_handler(int signum, siginfo_t *info, void *ucontext)
{
	uint64_t val;
	static unsigned count = 0;

	assert_return(signum > 0 && info);

	count++;
	if (count > 5) {
		dump_tree(cfghack);
		exit(EXIT_FAILURE);
	}

	switch (signum) {
	case SIGUSR1:
		debug(DBG_SIG, "Got a SIGUSR1");
		if (info->si_code != SI_ASYNCNL || info->si_signo != SIGUSR1 || !info->si_ptr) {
			debug(DBG_SIG, "unexpected values in siginfo");
			return;
		}
		debug(DBG_SIG, "SIGUSR1 struct dns_async: %p", info->si_ptr);
		hack_dns = info->si_ptr;
		val = 10000;
		break;
	case SIGINT:
		debug(DBG_SIG, "Got a SIGINT");
		val = 1000;
		break;
	case SIGHUP:
		debug(DBG_SIG, "Got a SIGHUP");
		val = 1000;
		break;
	case SIGTERM:
		debug(DBG_SIG, "Got a SIGTERM");
		val = 1;
		break;
	default:
		error("Got an unknown sig (%i)", signum);
		val = 1;
		break;
	}

	write(hack_efd, &val, sizeof(val));
}

static void
signalfd_init(struct cfg *cfg)
{
	int sfd;
	//sigset_t mask;
	struct signalfd_ev *sev;

	assert_return(cfg);

	sev = zmalloc(sizeof(*sev));
	if (!sev)
		die("malloc: %m");

	/*
	sigfillset(&mask);
	if (sigprocmask(SIG_BLOCK, &mask, NULL) < 0)
		die("sigprocmask: %m");

	sfd = signalfd(-1, &mask, SFD_CLOEXEC);
	if (sfd < 0)
		die("signalfd: %m");
	*/

	struct sigaction action;
	sigfillset(&action.sa_mask);
	action.sa_sigaction = hack_handler;
	action.sa_flags = SA_SIGINFO;

	sigaction(SIGINT, &action, NULL);
	sigaction(SIGHUP, &action, NULL);
	sigaction(SIGTERM, &action, NULL);
	sigaction(SIGUSR1, &action, NULL);

	action.sa_handler = SIG_IGN;
	action.sa_flags = 0;
	sigaction(SIGPIPE, &action, NULL);

	sfd = eventfd(0, EFD_CLOEXEC);
	if (sfd < 0)
		die("eventfd: %m");

	debug(DBG_SIG, "using fd %i", sfd);
	uring_task_init(cfg, &sev->task, "sev", uring_parent(cfg), signalfd_free);
	uring_task_set_fd(&sev->task, sfd);
	cfg->sev = sev;
	sev->cfg = cfg;
	hack_efd = sfd;
	uring_read(cfg, &sev->task, &sev->buf, sizeof(sev->buf), signalfd_read);
}

int
main(int argc, char **argv)
{
	struct cfg *cfg;
	struct server *server;
	unsigned server_count;

	cfg = cfg_init(argc, argv);
	cfghack = cfg;

	cfg_apply(cfg);

	cfg_read(cfg);

	uring_init(cfg);

	igmp_init(cfg);

	/* Drop CAP_NET_RAW (if we have it), only used for igmp */
	capng_clear(CAPNG_SELECT_BOTH);
	if (capng_update(CAPNG_ADD,
			 CAPNG_EFFECTIVE | CAPNG_PERMITTED,
			 CAP_NET_BIND_SERVICE))
		die("capng_update failed");

	if (capng_apply(CAPNG_SELECT_BOTH)) {
		/* Try clearing all caps, shouldn't fail */
		capng_clear(CAPNG_SELECT_BOTH);
		if (capng_apply(CAPNG_SELECT_BOTH))
			die("capng_apply failed");
	}

	signalfd_init(cfg);

	cfgdir_init(cfg);

	announce_init(cfg);

	announce_start(cfg->aev);

	idle_init(cfg);

	uring_task_put(cfg, &cfg->task);

	server_count = 0;
	list_for_each_entry(server, &cfg->servers, list)
		server_count++;

	sd_notifyf(0, "READY=1\n"
		   "STATUS=Running, %u server configurations loaded\n"
		   "MAINPID=%lu",
		   server_count,
		   (unsigned long)getpid());

	info("mcproxy started, %u server configurations loaded",
	     server_count);

	uring_event_loop(cfg);

	verbose("Exiting");

	xfree(cfg);

	if (debug_enabled(DBG_MALLOC))
		debug_resource_usage();

	fflush(stdout);
	fflush(stderr);
	exit(EXIT_SUCCESS);
}