#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>
#include <time.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] [--backed|-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");
	} 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 command 'help'.");
		return EXIT_FAILURE;
	}

	return EXIT_SUCCESS;
}

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

	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 command 'add-entry'.");
		return EXIT_FAILURE;
	} else if (argc > 1) {
		Log(LOG_LEVEL_ERROR, "Too many arguments supplied to command 'add-entry'.");
		return EXIT_FAILURE;
	}

	// Extract the entry identifier
	const char* entry_id = argv[0];

	// Create the entry
	Status status = AddEntry(entry_id);
	if (status != SUCCESS)
		return EXIT_FAILURE;

	Log(LOG_LEVEL_INFO, "Successfully added entry '%s'.", entry_id);
	return EXIT_SUCCESS;
}

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

	// Extract the entry identifier
	const char* entry_id = argv[0];

	// Remove the entry
	Status status = RemoveEntry(entry_id);
	if (status != SUCCESS)
		return EXIT_FAILURE;

	Log(LOG_LEVEL_INFO, "Successfully removed entry '%s'.", entry_id);
	return EXIT_SUCCESS;
}

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

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

	// Calculate the maximum length of the entries
	size_t max_length = 0;
	for (size_t i = 0; entries[i] != NULL; i++) {
		size_t length = strlen(entries[i]);
		if (length > max_length)
			max_length = length;
	}

	// Print the entries
	for (size_t i = 0; entries[i] != NULL; i++) {
		bool has_disk;
		status = DoesEntryDiskExist(entries[i], &has_disk);
		if (status != SUCCESS)
			return EXIT_FAILURE;

		if (has_disk) {
			DiskInfo info;
			status = GetEntryDiskInfo(entries[i], &info);
			if (status != SUCCESS)
				continue;

			// Format the size
			char* size_fmt;
			status = FormatSize(info.allocated, &size_fmt);
			if (status != SUCCESS)
				continue;

			fprintf(stdout, "%zu | %-*s | %-10s | %s%s\n", i, (int)max_length, entries[i], size_fmt, info.backing_identifier == NULL ? "No backing" : "Backed by ", info.backing_identifier == NULL ? "" : info.backing_identifier);
		} else
			fprintf(stdout, "%zu | %-*s | %s\n", i, (int)max_length, entries[i], "No disk ");
	}

	// Free the entries
	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 command 'clear-entries'.");
		return EXIT_FAILURE;
	}

	// Clear the entries
	Status status = ClearEntries();
	if (status != SUCCESS)
		return EXIT_FAILURE;

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

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

	// Extract the entry identifier
	const char* entry_id = argv[0];

	// Extract the options
	bool root = false;
	const char* size_str = NULL;

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

	for (int i = 1; i < argc; i++) {
		if (strcmp(argv[i], "--root") == 0 || strcmp(argv[i], "-r") == 0) {
			// Check if the option is duplicated
			if (root) {
				Log(LOG_LEVEL_ERROR, "Duplicate option '--root'.");
				return EXIT_FAILURE;
			}

			// Check if the option is the last argument
			if (i + 1 == argc) {
				Log(LOG_LEVEL_ERROR, "Missing argument for option '--root <size>'.");
				return EXIT_FAILURE;
			}

			// Extract the size
			size_str = argv[i + 1];
			root = true;

			// Skip the next argument as it is the size
			i++;
		} else if (strcmp(argv[i], "--backed") == 0 || strcmp(argv[i], "-b") == 0) {
			// Check if the option is duplicated
			if (backed) {
				Log(LOG_LEVEL_ERROR, "Duplicate option '--backed'.");
				return EXIT_FAILURE;
			}

			// Check if the option is the last argument
			if (i + 1 == argc) {
				Log(LOG_LEVEL_ERROR, "Missing argument for option '--backed <backing id>'.");
				return EXIT_FAILURE;
			}

			// Extract the backing identifier
			backing_id = argv[i + 1];
			backed = true;

			// Skip the next argument as it is the backing identifier
			i++;
		} else {
			Log(LOG_LEVEL_ERROR, "Unknown option '%s'.", argv[i]);
			return EXIT_FAILURE;
		}
	}

	// Don't allow both root and backed options
	if (root && backed) {
		Log(LOG_LEVEL_ERROR, "Cannot use both '--root' and '--backed' options.");
		return EXIT_FAILURE;
	}

	if (root) {
		// Parse the size
		uint64_t size;
		Status status = ParseSize(size_str, &size);
		if (status != SUCCESS)
			return EXIT_FAILURE;

		// Add the root disk
		status = AddRootEntryDisk(entry_id, size);
		if (status != SUCCESS)
			return EXIT_FAILURE;

		// Format the size
		char* size_fmt;
		status = FormatSize(size, &size_fmt);

		Log(LOG_LEVEL_INFO, "Successfully added root disk to entry '%s' of size %s.", entry_id, size_fmt);

		free(size_fmt);
		return EXIT_SUCCESS;
	} else if (backed) {
		// Check if the backing exists
		Status status = AddBackedEntryDisk(entry_id, backing_id);
		if (status != SUCCESS)
			return EXIT_FAILURE;

		Log(LOG_LEVEL_INFO, "Successfully added backed disk to entry '%s' with backing '%s'.", entry_id, backing_id);
		return EXIT_SUCCESS;
	} else {
		// Get the latest backing to use as the default backing
		char* backing_id;
		Status status = GetLatestBacking(&backing_id);
		if (status != SUCCESS) {
			Log(LOG_LEVEL_ERROR, "Failed to get the latest backing to use as the default backing.");
			return EXIT_FAILURE;
		}

		// Add the backed disk
		status = AddBackedEntryDisk(entry_id, backing_id);
		if (status != SUCCESS) {
			free(backing_id);
			return EXIT_FAILURE;
		}

		Log(LOG_LEVEL_INFO, "Successfully added backed disk to entry '%s' with backing '%s'.", entry_id, backing_id);

		free(backing_id);
		return EXIT_SUCCESS;
	}
}

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

	// Extract the entry identifier
	const char* entry_id = argv[0];

	// Remove the disk
	Status status = RemoveEntryDisk(entry_id);
	if (status != SUCCESS)
		return EXIT_FAILURE;

	Log(LOG_LEVEL_INFO, "Successfully removed disk from entry '%s'.", entry_id);
	return EXIT_SUCCESS;
}

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

	// Extract the entry identifier
	const char* entry_id = argv[0];

	// Extract the options
	bool update = false;
	bool backed = false;
	const char* backing_id = NULL;

	for (int i = 1; i < argc; i++) {
		if (strcmp(argv[i], "--update") == 0 || strcmp(argv[i], "-u") == 0) {
			// Check if the option is duplicated
			if (update) {
				Log(LOG_LEVEL_ERROR, "Duplicate option '--update'.");
				return EXIT_FAILURE;
			}

			update = true;
		} else if (strcmp(argv[i], "--backed") == 0 || strcmp(argv[i], "-b") == 0) {
			// Check if the option is duplicated
			if (backed) {
				Log(LOG_LEVEL_ERROR, "Duplicate option '--backed'.");
				return EXIT_FAILURE;
			}

			// Check if the option is the last argument
			if (i + 1 == argc) {
				Log(LOG_LEVEL_ERROR, "Missing argument for option '--backed <backing id>'.");
				return EXIT_FAILURE;
			}

			// Extract the backing identifier
			backing_id = argv[i + 1];
			backed = true;

			// Skip the next argument as it is the backing identifier
			i++;
		} else {
			Log(LOG_LEVEL_ERROR, "Unknown option '%s'.", argv[i]);
			return EXIT_FAILURE;
		}
	}

	// Don't allow both update and backed options
	if (update && backed) {
		Log(LOG_LEVEL_ERROR, "Cannot use both '--update' and '--backed' options.");
		return EXIT_FAILURE;
	}

	if (update) {
		char* backing_id;
		Status status = GetLatestBacking(&backing_id);
		if (status != SUCCESS) {
			Log(LOG_LEVEL_ERROR, "Failed to get the latest backing to use as the default backing.");
			return EXIT_FAILURE;
		}

		// Reset the disk
		status = ResetEntryDisk(entry_id, backing_id);
		if (status != SUCCESS) {
			free(backing_id);
			return EXIT_FAILURE;
		}

		Log(LOG_LEVEL_INFO, "Successfully reset disk from entry '%s' with backing '%s'.", entry_id, backing_id);

		free(backing_id);
		return EXIT_SUCCESS;
	} else if (backed) {
		// Reset the disk with the specified backing
		Status status = ResetEntryDisk(entry_id, backing_id);
		if (status != SUCCESS)
			return EXIT_FAILURE;

		Log(LOG_LEVEL_INFO, "Successfully reset disk from entry '%s' with backing '%s'.", entry_id, backing_id);
		return EXIT_SUCCESS;
	} else {
		// Reset the disk
		Status status = ResetEntryDisk(entry_id, NULL);
		if (status != SUCCESS)
			return EXIT_FAILURE;

		Log(LOG_LEVEL_INFO, "Successfully reset disk from entry '%s'.", entry_id);
		return EXIT_SUCCESS;
	}
}
int CommandTrimDisk(int argc, char* argv[]) {
	if (argc < 1) {
		Log(LOG_LEVEL_ERROR, "Too few arguments supplied to command 'trim-disk'.");
		return EXIT_FAILURE;
	} else if (argc > 1) {
		Log(LOG_LEVEL_ERROR, "Too many arguments supplied to command 'trim-disk'.");
		return EXIT_FAILURE;
	}

	// Extract the entry identifier
	const char* entry_id = argv[0];

	// Trim the disk
	Status status = TrimEntryDisk(entry_id);
	if (status != SUCCESS)
		return EXIT_FAILURE;

	Log(LOG_LEVEL_INFO, "Successfully trimmed disk from entry '%s'.", entry_id);
	return EXIT_SUCCESS;
}