#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <fcntl.h>
#include <string.h>
#include <sys/types.h>
#include <sys/wait.h>

#define POWER_STATE_F "/sys/power/state"
#define SLEEP_TIMEOUT 30

void
die(const char *msg)
{
	puts(msg);
	exit(1);
}

void
err(const char *msg)
{
	perror(msg);
	exit(1);
}

void
suspend()
{
	int fh = open(POWER_STATE_F, O_WRONLY);
	int n;
	const char *cmd = "mem";

	if (fh == -1)
		err("Opening "POWER_STATE_F);

	n = write(fh, cmd, 3);
	if (n == -1)
		err("Writing to "POWER_STATE_F);
}

int
main(int argc, char **argv)
{
	uid_t uid = getuid();
	gid_t gid = getgid();
	uid_t euid = geteuid();
	char **a = argv+1;
	char *lockProg = "slock";
	char *lockProgWl= "swaylock";
	pid_t pid;
	int ret;

	if (euid != 0)
		die("Program must be run as root/setuid");

	if (strcmp("wayland", getenv("XDG_SESSION_TYPE")) == 0)
		lockProg = lockProgWl;

	for (;*a;a++) {
		if (strcmp("-l", *a) == 0) {
			a++;
			if (*a)
				lockProg = *a;
			else
				die("-l must have an argument");
		} else {
			printf("Unknown argument: '%s'\n", *a);
			exit(1);
		}
	}

	pid = fork();
	switch (pid) {
	case -1:
		err("fork");
	case 0:
		ret = setuid(uid);
		if (ret == -1)
			err("setuid");
		ret = setgid(gid);
		if (ret == -1)
			err("setgid");

		execlp(lockProg, NULL);
		err("execlp");
	default:
		for(;;) {
			ret = 0;
			ret = waitpid(pid, NULL, WNOHANG);

			if (ret == 0) {
				suspend();
				sleep(SLEEP_TIMEOUT);
				continue;
			}

			puts("exiting");
			exit(0);
		}
	}

}