#include "disk.h"

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

result_t create_root_disk(const char* path, uint64_t size) {
	// Convert the size to a string
	char* size_str;
	result_t result = format(&size_str, "%llu", size);
	if (result != success())
		return result;

	// Create the disk
	int exit_code;
	// char* stdoutbuf;
	char* stderrbuf;

	result = execute(&exit_code, NULL, &stderrbuf, "qemu-img", "create", "-f", "qcow2", path, size_str, NULL);

	// Free the size string as it is no longer needed
	free(size_str);

	// Check for errors during the execution
	if (result != success())
		return result;

	// Check the exit code
	if (exit_code != 0) {
		// Remove all newlines from the stderr buffer
		for (char* c = stderrbuf; *c != '\0'; c++)
			if (*c == '\n')
				*c = ' ';

		result = failure("Failed to create root disk (%s).", stderrbuf);
		free(stderrbuf);
		return result;
	}

	// Free the stderr buffer
	free(stderrbuf);

	return success();
}

result_t create_backed_disk(const char* path, const char* backing_path) {
	// Create the disk
	int exit_code;
	// char* stdoutbuf;
	char* stderrbuf;

	result_t result = execute(&exit_code, NULL, &stderrbuf, "qemu-img", "create", "-f", "qcow2", "-F", "qcow2", "-b", backing_path, path, NULL);

	// Check for errors during the execution
	if (result != success())
		return result;

	// Check the exit code
	if (exit_code != 0) {
		// Remove all newlines from the stderr buffer
		for (char* c = stderrbuf; *c != '\0'; c++)
			if (*c == '\n')
				*c = ' ';

		result = failure("Failed to create backed disk (%s).", stderrbuf);
		free(stderrbuf);
		return result;
	}

	// Free the stderr buffer
	free(stderrbuf);

	return success();
}

result_t trim_disk(const char* path) {
	// Get the disk info
	disk_info_t info;
	result_t result = get_disk_info(&info, path);
	if (result != success())
		return result;

	// Generate a new path for the temporary disk
	char* temp_path;
	result = format(&temp_path, "%s.tmp", path);
	if (result != success()) {
		free_disk_info(&info);
		return result;
	}

	// Use convert to trim the disk
	int exit_code;
	// char* stdoutbuf;
	char* stderrbuf;

	if (info.backing_path == NULL)
		result = execute(&exit_code, NULL, &stderrbuf, "qemu-img", "convert", "-f", "qcow2", "-O", "qcow2", path, temp_path, NULL);
	else
		result = execute(&exit_code, NULL, &stderrbuf, "qemu-img", "convert", "-f", "qcow2", "-F", "qcow2", "-O", "qcow2", "-B", info.backing_path, path, temp_path, NULL);

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

	// Check for errors during the execution
	if (result != success()) {
		free(temp_path);
		return result;
	}

	// Check the exit code
	if (exit_code != 0) {
		// Remove all newlines from the stderr buffer
		for (char* c = stderrbuf; *c != '\0'; c++)
			if (*c == '\n')
				*c = ' ';

		result = failure("Failed to trim disk (%s).", stderrbuf);
		free(temp_path);
		free(stderrbuf);
		return result;
	}

	// Free the stderr buffer as it is no longer needed
	free(stderrbuf);

	// Replace the original disk with the temporary disk
	errno = 0;
	rename(temp_path, path);

	// Free the temporary path
	free(temp_path);

	// Check for errors during the renaming
	if (errno != 0) {
		// Try to remove the temporary disk if the renaming failed
		remove(temp_path);
		return failure("Failed to replace the original disk with the temporary disk (%s).", strerror(errno));
	}

	return success();
}

result_t reset_disk(const char* path) {
	// Get the disk info
	disk_info_t info;
	result_t result = get_disk_info(&info, path);
	if (result != success())
		return result;

	if (info.backing_path == NULL)
		// If the disk is not backed, simply create a new disk with the same size on top of it
		result = create_root_disk(path, info.size);
	else
		// If the disk is backed, create a new disk with the same backing path
		result = create_backed_disk(path, info.backing_path);

	// Free the disk info
	free_disk_info(&info);
	return result;
}

result_t reback_disk(const char* path, const char* backing_path) {
	// Create the disk
	int exit_code;
	// char* stdoutbuf;
	char* stderrbuf;

	result_t result = execute(&exit_code, NULL, &stderrbuf, "qemu-img", "rebase", "-u", "-f", "qcow2", "-F", "qcow2", "-b", backing_path, path, NULL);

	// Check for errors during the execution
	if (result != success())
		return result;

	// Check the exit code
	if (exit_code != 0) {
		// Remove all newlines from the stderr buffer
		for (char* c = stderrbuf; *c != '\0'; c++)
			if (*c == '\n')
				*c = ' ';

		result = failure("Failed to reback disk (%s).", stderrbuf);
		free(stderrbuf);
		return result;
	}

	// Free the stderr buffer
	free(stderrbuf);

	return success();
}

result_t get_disk_info(disk_info_t* _info, const char* path) {
	// Initialize the output parameters
	_info->size = 0;
	_info->allocated = 0;
	_info->backing_path = NULL;

	// Get the disk info
	int exit_code;
	char* stdoutbuf;
	char* stderrbuf;

	result_t result = execute(&exit_code, &stdoutbuf, &stderrbuf, "qemu-img", "info", "--output", "json", path, NULL);

	// Check for errors during the execution
	if (result != success())
		return result;

	// Check the exit code
	if (exit_code != 0) {
		// Remove all newlines from the stderr buffer
		for (char* c = stderrbuf; *c != '\0'; c++)
			if (*c == '\n')
				*c = ' ';

		result = failure("Failed to get disk info (%s).", stderrbuf);
		free(stdoutbuf);
		free(stderrbuf);
		return result;
	}

	// Free the stderr buffer as it is no longer needed
	free(stderrbuf);

	// Parse the JSON output
	json_object* root = json_tokener_parse(stdoutbuf);

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

	// Check for errors during the parsing
	if (root == NULL)
		return failure("Failed to parse the JSON output.");

	// Get the size
	json_object* size_obj;
	if (!json_object_object_get_ex(root, "virtual-size", &size_obj)) {
		json_object_put(root);
		return failure("Missing 'virtual-size' field in the JSON output. Cannot get the disk size.");
	}

	errno = 0;
	uint64_t size = json_object_get_int64(size_obj);
	if (errno != 0) {
		json_object_put(root);
		return failure("Failed to parse the 'virtual-size' field in the JSON output (%s).", strerror(errno));
	}

	// Get the allocated size
	json_object* allocated_obj;
	if (!json_object_object_get_ex(root, "actual-size", &allocated_obj)) {
		json_object_put(root);
		return failure("Missing 'actual-size' field in the JSON output. Cannot get the allocated size.");
	}

	errno = 0;
	uint64_t allocated = json_object_get_int64(allocated_obj);
	if (errno != 0) {
		json_object_put(root);
		return failure("Failed to parse the 'actual-size' field in the JSON output (%s).", strerror(errno));
	}

	// Get the backing file
	json_object* backing_obj;
	char* backing_path = NULL;
	if (json_object_object_get_ex(root, "backing-filename", &backing_obj)) {
		const char* backing_str = json_object_get_string(backing_obj);
		if (backing_str == NULL) {
			json_object_put(root);
			return failure("Failed to parse the 'backing-filename' field in the JSON output.");
		}

		// Duplicate the backing path string
		errno = 0;
		backing_path = strdup(backing_str);
		if (backing_path == NULL) {
			json_object_put(root);
			return failure("Failed to allocate memory for the backing path (%s).", strerror(errno));
		}
	}

	// Free the root object
	json_object_put(root);

	// Write the disk info to the output
	_info->size = size;
	_info->allocated = allocated;
	_info->backing_path = backing_path;
	return success();
}

void free_disk_info(disk_info_t* _info) {
	// Free the backing path string
	free(_info->backing_path);

	// Clear the fields
	_info->size = 0;
	_info->allocated = 0;
	_info->backing_path = NULL;
}