#include "sandbox.h"

#include "entry.h"
#include "backing.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <pwd.h>
#include <unistd.h>

Command COMMANDS[] = {
	{CommandHelp, "help", "[command]", "Prints this help message.",
	 "TODO: Add details."},
	{CommandVersion, "version", NULL, "Prints the version of the program.",
	 "TODO: Add details."},
	{},
	{CommandAddEntry, "add-entry", "<entry id>", "Adds a new entry to the sandbox.",
	 "TODO: Add details."},
	{CommandRemoveEntry, "remove-entry", "<entry id>", "Removes an entry from the sandbox.",
	 "TODO: Add details."},
	{CommandListEntries, "list-entries", NULL, "Lists all the entries in the sandbox.",
	 "TODO: Add details."},
	{CommandClearEntries, "clear-entries", NULL, "Clears all the entries from the sandbox.",
	 "TODO: Add details."},
	{},
	{CommandAddDisk, "add-disk", "<entry id> [--root|-r <size>] [--backed|-b <backing id>]", "Adds a new disk to an entry.",
	 "TODO: Add details."},
	{CommandRemoveDisk, "remove-disk", "<entry id>", "Removes the disk from an entry.",
	 "TODO: Add details."},
	{CommandResetDisk, "reset-disk", "<entry id> [--update|-u] [--backing|-b <backing id>]", "Resets the disk of an entry.",
	 "TODO: Add details."},
	{CommandTrimDisk, "trim-disk", "<entry id>", "Trims the disk of an entry.",
	 "TODO: Add details."}};

int main(int argc, char* argv[]) {
	struct passwd* pw = getpwnam(SANDBOX_USER);
	if (pw == NULL) {
		Log(LOG_LEVEL_ERROR, "Failed to get the 'sandbox' user (%s).", strerror(errno));
		return EXIT_FAILURE;
	}

	// Check if the current user is root or the 'sandbox' user
	if (getuid() != 0 && getuid() != pw->pw_uid) {
		Log(LOG_LEVEL_ERROR, "You must be root or the 'sandbox' user to use this program.");
		return EXIT_FAILURE;
	}

	// Try and switch to the 'sandbox' user if we are root
	if (geteuid() == 0) {
		if (setregid(pw->pw_gid, pw->pw_gid) == -1) {
			Log(LOG_LEVEL_ERROR, "Failed to set the real and effective group ID to the 'sandbox' user (%s).", strerror(errno));
			return EXIT_FAILURE;
		}

		if (setreuid(pw->pw_uid, pw->pw_uid) == -1) {
			Log(LOG_LEVEL_ERROR, "Failed to set the real and effective user ID to the 'sandbox' user (%s).", strerror(errno));
			return EXIT_FAILURE;
		}
	}

	// If there are no arguments, print the help message
	if (argc == 1)
		return CommandHelp(0, NULL);

	const char* input = argv[1];
	size_t input_length = strlen(input);

	// Try and find the best matching command
	const Command* command = NULL;
	for (size_t i = 0; i < sizeof(COMMANDS) / sizeof(COMMANDS[0]); i++) {
		if (COMMANDS[i].name == NULL)
			continue;

		if (input_length > strlen(COMMANDS[i].name))
			continue;

		if (strncmp(input, COMMANDS[i].name, input_length) == 0) {
			// If we have already found a matching command, then the input is ambiguous
			if (command != NULL) {
				Log(LOG_LEVEL_ERROR, "Ambiguous command '%s'.", input);
				return EXIT_FAILURE;
			}

			command = &COMMANDS[i];
		}
	}

	if (command == NULL) {
		Log(LOG_LEVEL_ERROR, "Unknown command '%s'.", input);
		return EXIT_FAILURE;
	}

	return command->handler(argc - 2, argv + 2);
}

int CommandHelp(int argc, char* argv[]) {
	// If there are no arguments, print the general help message
	if (argc == 0) {
		fprintf(stdout, "Usage: sandbox <command> [arguments] [options]\n");
		fprintf(stdout, "\n");
		fprintf(stdout, "Commands:\n");
		for (size_t i = 0; i < sizeof(COMMANDS) / sizeof(COMMANDS[0]); i++) {
			if (COMMANDS[i].name == NULL) {
				fprintf(stdout, "\n");
				continue;
			}

			fprintf(stdout, "    %s", COMMANDS[i].name);
			if (COMMANDS[i].arguments != NULL)
				fprintf(stdout, " %s", COMMANDS[i].arguments);
			fprintf(stdout, " - %s\n", COMMANDS[i].description);
		}

		fprintf(stdout, "\nFor more information, run 'sandbox help <command>'.\n");

		return EXIT_SUCCESS;
	} else if (argc == 1) { // If there is one argument, print the help message for the command
		const char* input = argv[0];

		for (size_t i = 0; i < sizeof(COMMANDS) / sizeof(COMMANDS[0]); i++) {
			if (COMMANDS[i].name == NULL)
				continue;

			if (strcmp(input, COMMANDS[i].name) == 0) {
				fprintf(stdout, "Usage: sandbox %s", COMMANDS[i].name);
				if (COMMANDS[i].arguments != NULL)
					fprintf(stdout, " %s", COMMANDS[i].arguments);
				fprintf(stdout, "\n    %s\n\n    %s\n", COMMANDS[i].description, COMMANDS[i].details);

				return EXIT_SUCCESS;
			}
		}

		Log(LOG_LEVEL_ERROR, "Unknown command '%s'.", input);
		return EXIT_FAILURE;
	} else { // If there are too many arguments, print an error message
		Log(LOG_LEVEL_ERROR, "Too many arguments supplied to 'help'.");
		return EXIT_FAILURE;
	}

	return EXIT_SUCCESS;
}

int CommandVersion(int argc, char* argv[]) {
	fprintf(stdout, "Sandbox utility v%s\n", VERSION);
	return EXIT_SUCCESS;
}

int CommandAddEntry(int argc, char* argv[]) {
	if (argc < 1) {
		Log(LOG_LEVEL_ERROR, "Too few arguments supplied to 'add-entry'.");
		return EXIT_FAILURE;
	} else if (argc > 1) {
		Log(LOG_LEVEL_ERROR, "Too many arguments supplied to 'add-entry'.");
		return EXIT_FAILURE;
	}

	char* entry_identifier = argv[0];

	Status status = AddEntry(entry_identifier);
	if (status != SUCCESS)
		return EXIT_FAILURE;

	Log(LOG_LEVEL_INFO, "Added entry '%s'.", entry_identifier);
	return EXIT_SUCCESS;
}

int CommandRemoveEntry(int argc, char* argv[]) {
	if (argc < 1) {
		Log(LOG_LEVEL_ERROR, "Too few arguments supplied to 'remove-entry'.");
		return EXIT_FAILURE;
	} else if (argc > 1) {
		Log(LOG_LEVEL_ERROR, "Too many arguments supplied to 'remove-entry'.");
		return EXIT_FAILURE;
	}

	char* entry_identifier = argv[0];

	Status status = RemoveEntry(entry_identifier);
	if (status != SUCCESS)
		return EXIT_FAILURE;

	Log(LOG_LEVEL_INFO, "Removed entry '%s'.", entry_identifier);
	return EXIT_SUCCESS;
}

int CommandListEntries(int argc, char* argv[]) {
	if (argc > 0) {
		Log(LOG_LEVEL_ERROR, "Too many arguments supplied to 'list-entries'.");
		return EXIT_FAILURE;
	}

	char** entries = NULL;
	Status status = ListEntries(&entries);
	if (status != SUCCESS)
		return EXIT_FAILURE;

	for (size_t i = 0; entries[i] != NULL; i++)
		fprintf(stdout, "%s\n", entries[i]);

	for (size_t i = 0; entries[i] != NULL; i++)
		free(entries[i]);
	free(entries);

	return EXIT_SUCCESS;
}

int CommandClearEntries(int argc, char* argv[]) {
	if (argc > 0) {
		Log(LOG_LEVEL_ERROR, "Too many arguments supplied to 'clear-entries'.");
		return EXIT_FAILURE;
	}

	Status status = ClearEntries();
	if (status != SUCCESS)
		return EXIT_FAILURE;

	Log(LOG_LEVEL_INFO, "Cleared all entries.");
	return EXIT_SUCCESS;
}

int CommandAddDisk(int argc, char* argv[]) {
	if (argc < 1) {
		Log(LOG_LEVEL_ERROR, "Too few arguments supplied to 'add-disk'.");
		return EXIT_FAILURE;
	}

	char* entry_identifier = argv[0];

	bool root = false;
	uint64_t size = 0;

	bool backed = false;
	const char* backing_identifier = NULL;

	for (int i = 1; i < argc; i++) {
		if (strcmp(argv[i], "--root") == 0 || strcmp(argv[i], "-r") == 0) {
			if (root) {
				Log(LOG_LEVEL_ERROR, "Duplicate '--root' argument supplied to 'add-disk'.");
				return EXIT_FAILURE;
			}

			if (i + 1 >= argc) {
				Log(LOG_LEVEL_ERROR, "Too few arguments supplied to 'add-disk'. Missing size for '--root <size>'.");
				return EXIT_FAILURE;
			}

			root = true;
			Status status = ParseSize(argv[i + 1], &size);
			if (status != SUCCESS)
				return EXIT_FAILURE;

			if (size == 0) {
				Log(LOG_LEVEL_ERROR, "Invalid size '%s' supplied to 'add-disk'.", argv[i + 1]);
				return EXIT_FAILURE;
			}

			i++; // Consumes the next argument (the size)
		} else if (strcmp(argv[i], "--backed") == 0 || strcmp(argv[i], "-b") == 0) {
			if (backed) {
				Log(LOG_LEVEL_ERROR, "Duplicate '--backed' argument supplied to 'add-disk'.");
				return EXIT_FAILURE;
			}

			if (i + 1 >= argc) {
				Log(LOG_LEVEL_ERROR, "Too few arguments supplied to 'add-disk'. Missing backing identifier for '--backed <backing id>'.");
				return EXIT_FAILURE;
			}

			backed = true;
			backing_identifier = argv[i + 1];

			if (!IsBackingIdentifierValid(backing_identifier)) {
				Log(LOG_LEVEL_ERROR, "Invalid backing identifier '%s' supplied to 'add-disk'.", backing_identifier);
				return EXIT_FAILURE;
			}

			i++; // Consumes the next argument (the backing identifier)
		} else {
			Log(LOG_LEVEL_ERROR, "Unknown argument '%s' supplied to 'add-disk'.", argv[i]);
			return EXIT_FAILURE;
		}
	}

	if (root && backed) {
		Log(LOG_LEVEL_ERROR, "The options '--root' and '--backed' cannot be used together in 'add-disk'.");
		return EXIT_FAILURE;
	}

	if (root) {
		Status status = AddRootEntryDisk(entry_identifier, size);
		if (status != SUCCESS)
			return EXIT_FAILURE;

		char* size_str = NULL;
		status = FormatSize(size, &size_str);
		Log(LOG_LEVEL_INFO, "Added root disk to entry '%s' with size %s.", entry_identifier, size_str);
		free(size_str);
	} else if (backed) {
		Status status = AddBackedEntryDisk(entry_identifier, backing_identifier);
		if (status != SUCCESS)
			return EXIT_FAILURE;

		Log(LOG_LEVEL_INFO, "Added backed disk to entry '%s' with backing '%s'.", entry_identifier, backing_identifier);
	} else {
		char* backing_identifier = NULL;
		Status status = GetLatestBacking(&backing_identifier);
		if (status != SUCCESS)
			return EXIT_FAILURE;

		status = AddBackedEntryDisk(entry_identifier, backing_identifier);
		if (status != SUCCESS) {
			free(backing_identifier);
			return EXIT_FAILURE;
		}

		Log(LOG_LEVEL_INFO, "Added backed disk to entry '%s' with backing '%s'.", entry_identifier, backing_identifier);
		free(backing_identifier);
	}

	return EXIT_SUCCESS;
}

int CommandRemoveDisk(int argc, char* argv[]) {
	if (argc < 1) {
		Log(LOG_LEVEL_ERROR, "Too few arguments supplied to 'remove-disk'.");
		return EXIT_FAILURE;
	} else if (argc > 1) {
		Log(LOG_LEVEL_ERROR, "Too many arguments supplied to 'remove-disk'.");
		return EXIT_FAILURE;
	}

	char* entry_identifier = argv[0];

	Status status = RemoveEntryDisk(entry_identifier);
	if (status != SUCCESS)
		return EXIT_FAILURE;

	Log(LOG_LEVEL_INFO, "Removed disk from entry '%s'.", entry_identifier);
	return EXIT_SUCCESS;
}

int CommandResetDisk(int argc, char* argv[]) {
	return 0;
}

int CommandTrimDisk(int argc, char* argv[]) {
	if (argc < 1) {
		Log(LOG_LEVEL_ERROR, "Too few arguments supplied to 'trim-disk'.");
		return EXIT_FAILURE;
	} else if (argc > 1) {
		Log(LOG_LEVEL_ERROR, "Too many arguments supplied to 'trim-disk'.");
		return EXIT_FAILURE;
	}

	char* entry_identifier = argv[0];

	Status status = TrimEntryDisk(entry_identifier);
	if (status != SUCCESS)
		return EXIT_FAILURE;

	Log(LOG_LEVEL_INFO, "Trimmed disk of entry '%s'.", entry_identifier);
	return EXIT_SUCCESS;
}