#include "disk.h"

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

Result CreateRootDisk(const char* disk_path, uint64_t size) {
	// Convert the size to a string
	char* size_str;
	if (Format(&size_str, "%lu", size) == FAILURE)
		return FAILURE;

	// Create the root disk
	int exit_code;
	// char* stdoutb; No need to capture stdout
	char* stderrb;

	Result result = RunExecutable(&exit_code, NULL, &stderrb, "qemu-img", "create", "-f", "qcow2", disk_path, size_str, NULL);

	// Free the size string
	free(size_str);

	// If the execution failed, return the error
	if (result != SUCCESS)
		return result;

	// If the exit code is non-zero, return the error
	if (exit_code != 0) {
		// If the error message is empty, return a generic error
		if (stderrb == NULL)
			Log(LOG_LEVEL_ERROR, "Failed to create root disk '%s'.", disk_path);
		else {
			// Remove all newlines from the error message
			for (char* c = stderrb; *c != '\0'; c++)
				if (*c == '\n')
					*c = ' ';

			Log(LOG_LEVEL_ERROR, "Failed to create root disk '%s' (%s).", disk_path, stderrb);
		}

		// Free the error message
		free(stderrb);
		return FAILURE;
	}

	// As there can still be an error message for a successful exit code, we need to free it
	free(stderrb);
	return SUCCESS;
}

Result CreateBackedDisk(const char* disk_path, const char* backing_disk_path) {
	// Create the backed disk
	int exit_code;
	// char* stdoutb; No need to capture stdout
	char* stderrb;

	Result result = RunExecutable(&exit_code, NULL, &stderrb, "qemu-img", "create", "-f", "qcow2", "-F", "qcow2", "-b", backing_disk_path, disk_path, NULL);

	// If the execution failed, return the error
	if (result != SUCCESS)
		return result;

	// If the exit code is non-zero, return the error
	if (exit_code != 0) {
		// If the error message is empty, return a generic error
		if (stderrb == NULL)
			Log(LOG_LEVEL_ERROR, "Failed to create backed disk '%s' with backing disk '%s'.", disk_path, backing_disk_path);
		else {
			// Remove all newlines from the error message
			for (char* c = stderrb; *c != '\0'; c++)
				if (*c == '\n')
					*c = ' ';

			Log(LOG_LEVEL_ERROR, "Failed to create backed disk '%s' with backing disk '%s' (%s).", disk_path, backing_disk_path, stderrb);
		}

		// Free the error message
		free(stderrb);
		return FAILURE;
	}

	// As there can still be an error message for a successful exit code, we need to free it
	free(stderrb);
	return SUCCESS;
}

Result TrimDisk(const char* disk_path) {
	// Get a temporary file path to store the trimmed disk
	char* trimmed_disk_path;
	if (Format(&trimmed_disk_path, "%s.trimmed", disk_path) == FAILURE)
		return FAILURE;

	// Get info about the disk
	DiskInfo info;
	if (GetDiskInfo(disk_path, &info) == FAILURE) {
		free(trimmed_disk_path);
		return FAILURE;
	}

	// Trim the disk
	int exit_code;
	// char* stdoutb; No need to capture stdout
	char* stderrb;

	// If the disk is not backed, we can just use the "convert" command to trim it
	Result result;

	if (info.backing_file_path == NULL)
		result = RunExecutable(&exit_code, NULL, &stderrb, "qemu-img", "convert", "-f", "qcow2", "-O", "qcow2", disk_path, trimmed_disk_path, NULL);
	else {
		// We need to get the option to specify the backing file
		char* backing_option;
		if (Format(&backing_option, "backing_file=%s", info.backing_file_path) == FAILURE) {
			free(trimmed_disk_path);
			FreeDiskInfo(&info);
			return FAILURE;
		}

		// Trim the disk using the "convert" command
		result = RunExecutable(&exit_code, NULL, &stderrb, "qemu-img", "convert", "-f", "qcow2", "-F", "qcow2", "-O", "qcow2", "-o", backing_option, disk_path, trimmed_disk_path, NULL);

		// Free the backing option
		free(backing_option);
	}

	// Free the disk info
	FreeDiskInfo(&info);

	// If the execution failed, return the error
	if (result != SUCCESS) {
		free(trimmed_disk_path);
		return result;
	}

	// If the exit code is non-zero, return the error
	if (exit_code != 0) {
		// If the error message is empty, return a generic error
		if (stderrb == NULL)
			Log(LOG_LEVEL_ERROR, "Failed to trim disk '%s'.", disk_path);
		else {
			// Remove all newlines from the error message
			for (char* c = stderrb; *c != '\0'; c++)
				if (*c == '\n')
					*c = ' ';

			Log(LOG_LEVEL_ERROR, "Failed to trim disk '%s' (%s).", disk_path, stderrb);
		}

		// Free the error message
		free(stderrb);
		free(trimmed_disk_path);
		return FAILURE;
	}

	// As there can still be an error message for a successful exit code, we need to free it
	free(stderrb);

	// Move the trimmed disk to the original path
	if (rename(trimmed_disk_path, disk_path) != 0) {
		Log(LOG_LEVEL_ERROR, "Failed to move trimmed disk '%s' to original path '%s' (%s).", trimmed_disk_path, disk_path, strerror(errno));
		unlink(trimmed_disk_path);
		free(trimmed_disk_path);
		return FAILURE;
	}

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

Result RebackDisk(const char* disk_path, const char* backing_disk_path) {
	// Reback the disk
	int exit_code;
	// char* stdoutb; No need to capture stdout
	char* stderrb;

	Result result = RunExecutable(&exit_code, NULL, &stderrb, "qemu-img", "rebase", "-f", "qcow2", "-F", "qcow2", "-u", "-b", backing_disk_path, disk_path, NULL);

	// If the execution failed, return the error
	if (result != SUCCESS)
		return result;

	// If the exit code is non-zero, return the error
	if (exit_code != 0) {
		// If the error message is empty, return a generic error
		if (stderrb == NULL)
			Log(LOG_LEVEL_ERROR, "Failed to reback disk '%s' with backing disk '%s'.", disk_path, backing_disk_path);
		else {
			// Remove all newlines from the error message
			for (char* c = stderrb; *c != '\0'; c++)
				if (*c == '\n')
					*c = ' ';

			Log(LOG_LEVEL_ERROR, "Failed to reback disk '%s' with backing disk '%s' (%s).", disk_path, backing_disk_path, stderrb);
		}

		// Free the error message
		free(stderrb);
		return FAILURE;
	}

	// As there can still be an error message for a successful exit code, we need to free it
	free(stderrb);
	return SUCCESS;
}

Result GetDiskInfo(const char* disk_path, DiskInfo* _info) {
	// Initialize the disk info
	*_info = (DiskInfo){0};

	// Get the disk info
	int exit_code;
	char* stdoutb;
	char* stderrb;

	Result result = RunExecutable(&exit_code, &stdoutb, &stderrb, "qemu-img", "info", "--output", "json", disk_path, NULL);

	// If the execution failed, return the error
	if (result != SUCCESS)
		return result;

	// If the exit code is non-zero, return the error
	if (exit_code != 0) {
		// If the error message is empty, return a generic error
		if (stderrb == NULL)
			Log(LOG_LEVEL_ERROR, "Failed to get info for disk '%s'.", disk_path);
		else {
			// Remove all newlines from the error message
			for (char* c = stderrb; *c != '\0'; c++)
				if (*c == '\n')
					*c = ' ';

			Log(LOG_LEVEL_ERROR, "Failed to get info for disk '%s' (%s).", disk_path, stderrb);
		}

		// Free the error message
		free(stdoutb);
		free(stderrb);
		return FAILURE;
	}

	// No need for the error message anymore
	free(stderrb);

	// Parse the JSON output
	json_object* root = json_tokener_parse(stdoutb);
	if (root == NULL) {
		Log(LOG_LEVEL_ERROR, "Failed to parse JSON output for disk '%s'.", disk_path);
		free(stdoutb);
		return FAILURE;
	}

	// No need for the JSON string anymore
	free(stdoutb);

	// Get the size of the disk
	json_object* size_obj;
	if (!json_object_object_get_ex(root, "virtual-size", &size_obj)) {
		Log(LOG_LEVEL_ERROR, "Failed to get virtual size for disk '%s'.", disk_path);
		json_object_put(root);
		FreeDiskInfo(_info);
		return FAILURE;
	}

	errno = 0;
	_info->size = json_object_get_int64(size_obj);
	if (errno != 0) {
		Log(LOG_LEVEL_ERROR, "Failed to parse virtual size for disk '%s'.", disk_path);
		json_object_put(root);
		FreeDiskInfo(_info);
		return FAILURE;
	}

	// Get the allocated size of the disk
	json_object* allocation_obj;
	if (!json_object_object_get_ex(root, "actual-size", &allocation_obj)) {
		Log(LOG_LEVEL_ERROR, "Failed to get actual size for disk '%s'.", disk_path);
		json_object_put(root);
		FreeDiskInfo(_info);
		return FAILURE;
	}

	errno = 0;
	_info->allocated = json_object_get_int64(allocation_obj);
	if (errno != 0) {
		Log(LOG_LEVEL_ERROR, "Failed to parse actual size for disk '%s'.", disk_path);
		json_object_put(root);
		FreeDiskInfo(_info);
		return FAILURE;
	}

	// Get the backing file of the disk
	json_object* backing_file_obj;
	if (!json_object_object_get_ex(root, "backing-filename", &backing_file_obj))
		_info->backing_file_path = NULL;
	else {
		const char* backing_file_path = json_object_get_string(backing_file_obj);
		if (backing_file_path == NULL) {
			Log(LOG_LEVEL_ERROR, "Failed to parse backing file for disk '%s'.", disk_path);
			json_object_put(root);
			FreeDiskInfo(_info);
			return FAILURE;
		}

		if (Duplicate(backing_file_path, &_info->backing_file_path) == FAILURE) {
			json_object_put(root);
			FreeDiskInfo(_info);
			return FAILURE;
		}

		// Get the backing identifier of the disk (basename)
		if (Duplicate(basename(_info->backing_file_path), &_info->backing_identifier) == FAILURE) {
			json_object_put(root);
			FreeDiskInfo(_info);
			return FAILURE;
		}
	}

	// Free the JSON object
	json_object_put(root);
	return SUCCESS;
}

Result FreeDiskInfo(DiskInfo* info) {
	free(info->backing_file_path);
	free(info->backing_identifier);
	return SUCCESS;
}