#include "disk.h"

#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <sys/stat.h>
#include <stdio.h>
#include <json-c/json.h>

Result create_root_disk(const char* path, uint64_t size, mode_t permissions) {
	// Convert the size to a string
	char* size_str;
	Result result = format(&size_str, "%lu", size);
	if (result != SUCCESS)
		return result;

	// Create the disk
	int exit_code;
	char* stderr_buffer = NULL;
	result = run_executable(&exit_code, NULL, &stderr_buffer, "qemu-img", "create", "-f", "qcow2", path, size_str, NULL);

	// Free the size string
	free(size_str);

	if (result != SUCCESS) {
		free(stderr_buffer);
		return result;
	}

	// Check the exit code
	if (exit_code != EXIT_SUCCESS) {
		if (stderr_buffer == NULL)
			log_message(LOG_LEVEL_ERROR, "Failed to create the disk '%s'.", path);
		else {
			// Remove newlines from the error message
			for (size_t i = 0; i < strlen(stderr_buffer); i++)
				if (stderr_buffer[i] == '\n')
					stderr_buffer[i] = ' ';

			log_message(LOG_LEVEL_ERROR, "Failed to create the disk '%s' (%s).", path, stderr_buffer);
			free(stderr_buffer);
		}

		return FAILURE;
	}

	// Set the permissions of the disk
	if (chmod(path, permissions) != 0) {
		log_message(LOG_LEVEL_ERROR, "Failed to set the permissions of the disk '%s' (%s).", path, strerror(errno));

		// Try to remove the disk if the permissions could not be set
		unlink(path);

		return FAILURE;
	}

	return SUCCESS;
}

Result create_backed_disk(const char* path, const char* backing_disk, mode_t permissions) {
	// Create the disk
	int exit_code;
	char* stderr_buffer = NULL;
	Result result = run_executable(&exit_code, NULL, &stderr_buffer, "qemu-img", "create", "-f", "qcow2", "-F", "qcow2", "-b", backing_disk, path, NULL);

	if (result != SUCCESS) {
		free(stderr_buffer);
		return result;
	}

	// Check the exit code
	if (exit_code != EXIT_SUCCESS) {
		if (stderr_buffer == NULL)
			log_message(LOG_LEVEL_ERROR, "Failed to create the disk '%s'.", path);
		else {
			// Remove newlines from the error message
			for (size_t i = 0; i < strlen(stderr_buffer); i++)
				if (stderr_buffer[i] == '\n')
					stderr_buffer[i] = ' ';

			log_message(LOG_LEVEL_ERROR, "Failed to create the disk '%s' (%s).", path, stderr_buffer);
			free(stderr_buffer);
		}

		return FAILURE;
	}

	// Set the permissions of the disk
	if (chmod(path, permissions) != 0) {
		log_message(LOG_LEVEL_ERROR, "Failed to set the permissions of the disk '%s' (%s).", path, strerror(errno));

		// Try to remove the disk if the permissions could not be set
		unlink(path);

		return FAILURE;
	}

	return SUCCESS;
}

Result trim_disk(const char* path) {
	char* tmp_path;
	Result result = format(&tmp_path, "%s.tmp", path);
	if (result != SUCCESS)
		return result;

	DiskInfo info;
	result = get_disk_info(path, &info);
	if (result != SUCCESS) {
		free(tmp_path);
		return result;
	}

	// Create the trimmed disk
	int exit_code;
	char* stderr_buffer = NULL;

	if (info.backing_file == NULL) {
		// If the disk is not backed, we can simply convert it to a new disk
		result = run_executable(&exit_code, NULL, &stderr_buffer, "qemu-img", "convert", "-f", "qcow2", "-O", "qcow2", path, tmp_path, NULL);
	} else {
		// If the disk is backed, we need to specify the backing file
		char* backing_disk_arg;
		result = format(&backing_disk_arg, "backing_file=%s", info.backing_file);
		if (result != SUCCESS) {
			free(tmp_path);
			free_disk_info(&info);
			return result;
		}

		result = run_executable(&exit_code, NULL, &stderr_buffer, "qemu-img", "convert", "-f", "qcow2", "-F", "qcow2", "-O", "qcow2", "-o", backing_disk_arg, path, tmp_path, NULL);

		free(backing_disk_arg);
	}

	// Free the disk info as it is no longer needed
	free_disk_info(&info);

	if (result != SUCCESS) {
		free(tmp_path);
		free(stderr_buffer);
		return result;
	}

	// Check the exit code
	if (exit_code != EXIT_SUCCESS) {
		if (stderr_buffer == NULL)
			log_message(LOG_LEVEL_ERROR, "Failed to trim the disk '%s'.", path);
		else {
			// Remove newlines from the error message
			for (size_t i = 0; i < strlen(stderr_buffer); i++)
				if (stderr_buffer[i] == '\n')
					stderr_buffer[i] = ' ';

			log_message(LOG_LEVEL_ERROR, "Failed to trim the disk '%s' (%s).", path, stderr_buffer);
			free(stderr_buffer);
		}

		free(tmp_path);
		return FAILURE;
	}

	// Replace the original disk with the trimmed disk
	if (rename(tmp_path, path) != 0) {
		log_message(LOG_LEVEL_ERROR, "Failed to replace the disk '%s' with the trimmed disk '%s' (%s).", path, tmp_path, strerror(errno));

		unlink(tmp_path);
		free(tmp_path);
		return FAILURE;
	}

	// Free the temporary path
	free(tmp_path);
	return SUCCESS;
}

Result rebase_disk(const char* path, const char* backing_disk) {
	// Rebase the disk
	int exit_code;
	char* stderr_buffer = NULL;
	Result result = run_executable(&exit_code, NULL, &stderr_buffer, "qemu-img", "rebase", "-u", "-b", backing_disk, path, NULL);

	if (result != SUCCESS) {
		free(stderr_buffer);
		return result;
	}

	// Check the exit code
	if (exit_code != EXIT_SUCCESS) {
		if (stderr_buffer == NULL)
			log_message(LOG_LEVEL_ERROR, "Failed to rebase the disk '%s'.", path);
		else {
			// Remove newlines from the error message
			for (size_t i = 0; i < strlen(stderr_buffer); i++)
				if (stderr_buffer[i] == '\n')
					stderr_buffer[i] = ' ';

			log_message(LOG_LEVEL_ERROR, "Failed to rebase the disk '%s' (%s).", path, stderr_buffer);
			free(stderr_buffer);
		}

		return FAILURE;
	}

	return SUCCESS;
}

Result get_disk_info(const char* path, DiskInfo* out_info) {
	// Initialize the output
	out_info->backing_file = NULL;
	out_info->actual_size = 0;
	out_info->virtual_size = 0;

	// Get the information about the disk
	int exit_code;
	char* stdout_buffer = NULL;
	Result result = run_executable(&exit_code, &stdout_buffer, NULL, "qemu-img", "info", "--output", "json", path, NULL);
	if (result != SUCCESS) {
		free(stdout_buffer);
		return result;
	}

	// Check the exit code
	if (exit_code != EXIT_SUCCESS) {
		if (stdout_buffer == NULL)
			log_message(LOG_LEVEL_ERROR, "Failed to get information about the disk '%s'.", path);
		else {
			// Remove newlines from the error message
			for (size_t i = 0; i < strlen(stdout_buffer); i++)
				if (stdout_buffer[i] == '\n')
					stdout_buffer[i] = ' ';

			log_message(LOG_LEVEL_ERROR, "Failed to get information about the disk '%s' (%s).", path, stdout_buffer);
			free(stdout_buffer);
		}

		return FAILURE;
	}

	// Parse the JSON output
	json_object* root = json_tokener_parse(stdout_buffer);
	if (root == NULL) {
		log_message(LOG_LEVEL_ERROR, "Failed to parse the JSON output of 'qemu-img info'.");
		free(stdout_buffer);
		return FAILURE;
	}

	// Free the stdout buffer, as it is no longer needed
	free(stdout_buffer);

	// Get the virtual size
	json_object* virtual_size_obj;
	if (json_object_object_get_ex(root, "virtual-size", &virtual_size_obj))
		out_info->virtual_size = json_object_get_int64(virtual_size_obj);

	// Get the actual size
	json_object* actual_size_obj;
	if (json_object_object_get_ex(root, "actual-size", &actual_size_obj))
		out_info->actual_size = json_object_get_int64(actual_size_obj);

	// Get the backing file
	json_object* backing_file_obj;
	if (json_object_object_get_ex(root, "backing-filename", &backing_file_obj)) {
		out_info->backing_file = strdup(json_object_get_string(backing_file_obj));

		if (out_info->backing_file == NULL) {
			log_message(LOG_LEVEL_ERROR, "Failed to allocate memory for the backing file path.");
			json_object_put(root);
			return OUT_OF_MEMORY;
		}
	}

	// Free the JSON object
	json_object_put(root);

	return SUCCESS;
}

void free_disk_info(DiskInfo* info) {
	free(info->backing_file);

	info->backing_file = NULL;
	info->actual_size = 0;
	info->virtual_size = 0;
}