#include "../util.h"

#include <sys/param.h>
#include <sys/mount.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <signal.h>
#include <pwd.h>

static char *Root;

NAMED_ENUM(mount_type, layer, devfs, );

struct mount
{
	mount_type Type;
	union
	{
		struct
		{
			char *Source;
		} layer;
		struct
		{
			char *Target;
		} devfs;
	};
};

static struct mount *Mounts;
static size_t        MountsCount;
static bool          MountsCleanup;
static pid_t         SlavePID;
static fd            SlaveSocket;

static char const *Arg0;

static void __attribute__((noreturn)) usage(void)
{
	fprintf(stderr, "usage: %1$s [recipe-file] [name]\n", Arg0);
	exit(EX_USAGE);
}

NAMED_ENUM(message_type, run, chdir, user, );

#define MAX_MESSAGE_SIZE                 8192
#define MAX_SLAVE_COMMAND_ARGUMENT_COUNT 32

struct
{
	message_type Type;
	union
	{
		struct
		{
			size_t ArgCount;
			size_t Arg[MAX_SLAVE_COMMAND_ARGUMENT_COUNT];
			char   ArgData[0];
		} run;
		struct
		{
			char Path[0];
		} chdir;
		struct
		{
			char Username[0];
		} user;
	};
} typedef message;

static string_array
Split(char const *Line, size_t LineLen)
{
	string_array E;
	E._              = NULL;
	E.Count          = 0;
	bool   PartStart = true;
	size_t Mark      = 0;
	for(size_t i = 0; i < LineLen; ++i)
	{
		if(Line[i] == ' ')
		{
			if(PartStart == false)
			{
				PartStart        = true;
				E._[E.Count - 1] = strndup(Line + Mark, i - Mark);
			}
		} else if(PartStart == true)
		{
			PartStart = false;
			++E.Count;
			Mark = i;
			E._  = reallocarray(E._, E.Count, sizeof(*E._));
			if(E._ == NULL)
			{
				fprintf(stderr, "fatal: reallocarray failed\n");
				exit(-1);
			}
		}
	}
	if(PartStart == false)
	{
		PartStart        = true;
		E._[E.Count - 1] = strndup(Line + Mark, LineLen - Mark);
	}
	E._          = reallocarray(E._, E.Count + 1, sizeof(*E._));
	E._[E.Count] = NULL;
	return E;
}

static void
ExpandMounts(void)
{
	struct mount *NewMounts = reallocarray(Mounts, MountsCount + 1, sizeof(*NewMounts));
	if(NewMounts == NULL)
	{
		fprintf(stderr, "reallocarray failed: %s\n", strerror(errno));
		exit(-1);
	}
	Mounts = NewMounts;
}

static int
command_layer(string_array E)
{
	if(E.Count != 1)
	{
		fprintf(stderr, "command expects 1 argument\n");
		return -1;
	}
	char const *Source = E._[0];
	ExpandMounts();
	char Target[MAXPATHLEN];
	snprintf(Target, sizeof(Target), "%s", Root);
	if(mkdir(Target, 755) == -1)
	{
		if(errno != EEXIST)
		{
			fprintf(stderr, "mkdir %s: %s\n", Target, strerror(errno));
			return -1;
		}
	}
	int Result = mount_unionfs(Target, Source, true);
	if(Result == -1) return -1;
	Mounts[MountsCount++] = (struct mount){
		.Type  = mount_type_layer,
		.layer = { strdup(Source) },
	};
	return 0;
}

static int
command_devfs(string_array E)
{
	if(E.Count != 1)
	{
		fprintf(stderr, "command expects 1 argument\n");
		return -1;
	}
	char const *TargetRelative = E._[0];
	ExpandMounts();
	char Target[MAXPATHLEN];
	snprintf(Target, sizeof(Target), "%s/%s", Root, E._[0]);
	int Result = mount_devfs(Target);
	if(Result == -1) return -1;
	Mounts[MountsCount++] = (struct mount){
		.Type  = mount_type_devfs,
		.devfs = { strdup(TargetRelative) },
	};
	return 0;
}

static int
command_copy(string_array E)
{
	if(E.Count != 2)
	{
		fprintf(stderr, "command expects 2 arguments\n");
		return -1;
	}
	char const *Source = E._[0];
	char        Target[MAXPATHLEN];
	snprintf(Target, sizeof(Target), "%s/%s", Root, E._[1]);
	pid_t PID = fork();
	if(PID == -1) return -1;
	if(PID == 0)
	{
		MountsCleanup = false;
		close(STDIN_FILENO);
		dup2(STDERR_FILENO, STDOUT_FILENO);
		close_all();
		fprintf(stderr, "rsync --archive --xattrs %s %s\n", Source, Target);
		execlp("rsync", "rsync", "--archive", "--xattrs", Source, Target, NULL);
		exit(-1);
	}
	int Exit;
	waitpid(PID, &Exit, 0);
	if(!WIFEXITED(Exit) || (WEXITSTATUS(Exit) != 0)) return -1;
	return 0;
}

static int
command_run(string_array E)
{
	if(E.Count > MAX_SLAVE_COMMAND_ARGUMENT_COUNT)
	{
		fprintf(stderr, "too many arguments\n");
		return -1;
	}
	char     Buffer[MAX_MESSAGE_SIZE];
	message *M      = (message *)Buffer;
	M->Type         = message_type_run;
	M->run.ArgCount = E.Count;
	size_t ArgAt    = 0;
	size_t At       = 0;
	for(size_t i = 0; i < E.Count; ++i)
	{
		M->run.Arg[ArgAt] = At;
		size_t Length     = strlen(E._[i]);
		if(At + Length >= sizeof(Buffer) - OffsetOf(message, run.ArgData))
		{
			fprintf(stderr, "command too long\n");
			return -1;
		}
		memcpy(M->run.ArgData + At, E._[i], Length);
		M->run.ArgData[At + Length] = 0;
		At += Length + 1;
		++ArgAt;
	}
	send(SlaveSocket, M, OffsetOf(message, run.ArgData) + At, 0);
	int Status;
	if(recv(SlaveSocket, &Status, sizeof(Status), 0) == -1)
	{
		fprintf(stderr, "recv: %s\n", strerror(errno));
		return -1;
	}
	return Status;
}

static int
command_chdir(string_array E)
{
	if(E.Count != 1)
	{
		fprintf(stderr, "requires 1 argument\n");
		return -1;
	}
	char     Buffer[MAX_MESSAGE_SIZE];
	message *M    = (message *)Buffer;
	M->Type       = message_type_chdir;
	size_t Length = strlen(E._[0]);
	if(Length >= sizeof(Buffer) - OffsetOf(message, chdir.Path))
	{
		fprintf(stderr, "chdir path too long\n");
		return -1;
	}
	memcpy(M->chdir.Path, E._[0], Length);
	M->chdir.Path[Length] = 0;
	send(SlaveSocket, M, OffsetOf(message, chdir.Path) + Length + 1, 0);
	int Status;
	if(recv(SlaveSocket, &Status, sizeof(Status), 0) == -1)
	{
		fprintf(stderr, "recv: %s\n", strerror(errno));
		return -1;
	}
	return Status;
}

static int
command_user(string_array E)
{
	if(E.Count != 1)
	{
		fprintf(stderr, "requires 1 argument\n");
		return -1;
	}
	char     Buffer[MAX_MESSAGE_SIZE];
	message *M    = (message *)Buffer;
	M->Type       = message_type_user;
	size_t Length = strlen(E._[0]);
	if(Length >= sizeof(Buffer) - OffsetOf(message, user.Username))
	{
		fprintf(stderr, "username too long\n");
		return -1;
	}
	memcpy(M->user.Username, E._[0], Length);
	M->user.Username[Length] = 0;
	send(SlaveSocket, M, OffsetOf(message, user.Username) + Length + 1, 0);
	int Status;
	if(recv(SlaveSocket, &Status, sizeof(Status), 0) == -1)
	{
		fprintf(stderr, "recv: %s\n", strerror(errno));
		return -1;
	}
	return Status;
}

static int
Execute(char const *Line, size_t LineLen)
{
	string_array Args = Split(Line, LineLen);
	if(Args.Count == 0)
	{
		Free_string_array(&Args);
		return 0;
	}
	int          Result;
	char        *Command = Args._[0];
	string_array Rest    = (string_array){ ._ = Args._ + 1, .Count = Args.Count - 1 };
	if(strcmp(Command, "RUN") == 0)
	{
		Result = command_run(Rest);
	} else if(strcmp(Command, "LAYER") == 0)
	{
		Result = command_layer(Rest);
	} else if(strcmp(Command, "DEVFS") == 0)
	{
		Result = command_devfs(Rest);
	} else if(strcmp(Command, "CHDIR") == 0)
	{
		Result = command_chdir(Rest);
	} else if(strcmp(Command, "COPY") == 0)
	{
		Result = command_copy(Rest);
	} else if(strcmp(Command, "USER") == 0)
	{
		Result = command_user(Rest);
	} else
	{
		fprintf(stderr, "unrecognized command: %s\n", Command);
		Result = -1;
	}
	Free_string_array(&Args);
	return Result;
}

static int
unmount_single(struct mount *E)
{
	char const *Target;
	switch(E->Type)
	{
		case mount_type_layer: Target = Root; break;
		case mount_type_devfs:
		{
			char TargetBuffer[MAXPATHLEN];
			snprintf(TargetBuffer, sizeof(TargetBuffer), "%s/%s", Root, E->devfs.Target);
			Target = TargetBuffer;
			break;
		}
		case mount_type_COUNT:
		default: __builtin_unreachable();
	}
	while(unmount(Target, MNT_FORCE) == -1)
	{
		fprintf(stderr, "unmount %s failed: %s\n", Target, strerror(errno));
		if(errno != EBUSY) return -1;
		sleep(1);
	}
	return 0;
}

static void
unmount_all(void)
{
	for(size_t i = MountsCount - 1; i < MountsCount; --i)
	{
		unmount_single(Mounts + i);
		switch(Mounts[i].Type)
		{
			case mount_type_layer: free(Mounts[i].layer.Source); break;
			case mount_type_devfs: free(Mounts[i].devfs.Target); break;
			case mount_type_COUNT:
			default: __builtin_unreachable();
		}
	}
	if(Mounts != NULL) free(Mounts);
}

static int
slave_command_run(string_array E)
{
	if(E.Count == 0)
	{
		fprintf(stderr, "run what?\n");
		return -1;
	}
	pid_t PID = fork();
	if(PID == -1) return -1;
	if(PID == 0)
	{
		MountsCleanup = false;
		close(STDIN_FILENO);
		dup2(STDERR_FILENO, STDOUT_FILENO);
		close_all();
		execvp(E._[0], E._);
		fprintf(stderr, "failed to exec '%s': %s\n", E._[0], strerror(errno));
		exit(-1);
	}
	int Exit;
	waitpid(PID, &Exit, 0);
	if(!WIFEXITED(Exit) || (WEXITSTATUS(Exit) != 0)) return -1;
	return 0;
}

static int
chuser(char *Username)
{
	errno            = 0;
	struct passwd *E = getpwnam(Username);
	if(E == NULL)
	{
		if(errno != 0) fprintf(stderr, "getpwnam(%s): %s\n", Username, strerror(errno));
		else
			fprintf(stderr, "user %s does not exist\n", Username);
		return -1;
	}
	uid_t UID = E->pw_uid;
	gid_t GID = E->pw_gid;
	if(setgid(GID) == -1)
	{
		fprintf(stderr, "setgid(%u): %s\n", GID, strerror(errno));
		return -1;
	}
	if(setuid(UID) == -1)
	{
		fprintf(stderr, "setuid(%u): %s\n", UID, strerror(errno));
		return -1;
	}
	return 0;
}

static int
slave(fd FD)
{
	char     Buffer[MAX_MESSAGE_SIZE];
	message *M = (message *)Buffer;
	ssize_t  Size;
	while((Size = recv(FD, Buffer, sizeof(Buffer), 0)))
	{
		if(Size == -1)
		{
			if(errno == EINTR) continue;
			fprintf(stderr, "recv: %s\n", strerror(errno));
			break;
		}
		if(Size == 0) break;
		int Result;
		switch(M->Type)
		{
			case message_type_run:
			{
				string_array E;
				E._          = NULL;
				E.Count      = M->run.ArgCount;
				E._          = reallocarray(E._, E.Count + 1, sizeof(*E._));
				E._[E.Count] = NULL;
				for(size_t i = 0; i < M->run.ArgCount; ++i) E._[i] = M->run.ArgData + M->run.Arg[i];
				Result = slave_command_run(E);
				free(E._);
				break;
			}
			case message_type_chdir:
			{
				Result = chdir(M->chdir.Path);
				break;
			}
			case message_type_user:
			{
				Result = chuser(M->user.Username);
				break;
			}
			case message_type_COUNT:
			default: __builtin_unreachable();
		}
		send(FD, &Result, sizeof(Result), 0);
	}
	close(FD);
	return 0;
}

static void
StopSlave(void)
{
	kill(SlavePID, SIGINT);
	int Exit;
	waitpid(SlavePID, &Exit, 0);
}

static int
SpawnSlave(void)
{
	fd FD[2];
	if(socketpair(PF_UNIX, SOCK_DGRAM, 0, FD) == -1)
	{
		fprintf(stderr, "socketpair %s\n", strerror(errno));
		return -1;
	}
	pid_t PID = fork();
	if(PID == -1) return -1;
	if(PID == 0)
	{
		close(FD[1]);
		close(STDIN_FILENO);
		dup2(STDERR_FILENO, STDOUT_FILENO);
		if(chroot(Root) == -1)
		{
			fprintf(stderr, "chroot %s: %s\n", Root, strerror(errno));
			exit(-1);
		}
		if(chdir("/") == -1)
		{
			fprintf(stderr, "chdir /: %s\n", strerror(errno));
			exit(-1);
		}
		exit(slave(FD[0]));
	}
	close(FD[0]);
	SlaveSocket = FD[1];
	SlavePID    = PID;
	return 0;
}

static void
Cleanup(void)
{
	if(SlavePID > 0) StopSlave();
	if(MountsCleanup == true) unmount_all();
}

static int
Shell(void)
{
	pid_t PID = fork();
	if(PID == -1)
	{
		fprintf(stderr, "fork: %s\n", strerror(errno));
		return -1;
	}
	if(PID == 0)
	{
		close_all();
		if(chroot(Root) == -1)
		{
			fprintf(stderr, "chroot %s: %s\n", Root, strerror(errno));
			exit(-1);
		}
		if(chdir("/") == -1)
		{
			fprintf(stderr, "chdir /: %s\n", strerror(errno));
			exit(-1);
		}
		execlp("/bin/sh", "sh", NULL);
		fprintf(stderr, "execlp /bin/sh: %s\n", strerror(errno));
		exit(-1);
	}
	int Exit;
	waitpid(PID, &Exit, 0);
	if(!WIFEXITED(Exit) || (WEXITSTATUS(Exit) != 0)) return -1;
	return 0;
}

int
main(int ArgCount, char **Arg)
{
	Arg0 = Arg[0];
	--ArgCount;
	++Arg;
	if(ArgCount != 2) usage();
	char const *Path = Arg[0];
	Root             = Arg[1];
	ArgCount -= 2;
	Arg += 2;
	if(mkdir(Root, 0755) == -1)
	{
		fprintf(stderr, "mkdir %s: %s\n", Root, strerror(errno));
		exit(-1);
	}
	atexit(Cleanup);
	if(SpawnSlave() == -1) exit(-1);
	FILE *F = fopen(Path, "r");
	if(F == NULL)
	{
		fprintf(stderr, "fopen(\"%s\", \"r\")\n", Path);
		exit(-1);
	}
	MountsCleanup    = true;
	char   *Line     = NULL;
	size_t  LineSize = 0;
	ssize_t LineLength;
	size_t  LineAt = 0;
	while((LineLength = getline(&Line, &LineSize, F)))
	{
		if(LineLength <= 0) break;
		++LineAt;
		if(Line[LineLength - 1] == '\n') Line[--LineLength] = '\0';
		if(LineLength == 0) continue;
		if(Line[0] == '#') continue;
		if(Execute(Line, (size_t)LineLength) == -1)
		{
			fprintf(stderr, "(line %zu)\n", LineAt);
			fprintf(stderr, "starting a shell inside %s\n", Root);
			Shell();
			break;
		}
	}
	printf("LAYER %s\n", Root);
	for(size_t i = 0; i < MountsCount; ++i)
	{
		struct mount *E = Mounts + i;
		switch(Mounts[i].Type)
		{
			case mount_type_layer:
			{
				printf("LAYER %s\n", E->layer.Source);
				break;
			}
			case mount_type_devfs:
			{
				printf("DEVFS %s\n", E->devfs.Target);
				break;
			}
			case mount_type_COUNT:
			default: __builtin_unreachable();
		}
	}
	return 0;
}