#include "builder.h"
#include "commands.h"
#include "exe.h"
#include "image.h"
#include "substitute.h"

#include <blue/cmd.h>
#include <blue/io/file.h>
#include <blue/io/path.h>
#include <blue/object/dict.h>
#include <blue/object/number.h>
#include <blue/object/string.h>
#include <blue/serial.h>
#include <blue/term/print.h>
#include <errno.h>
#include <inttypes.h>
#include <stdio.h>
#include <stdlib.h>

enum {
	OPT_SPECIFICATION,
	OPT_SPECIFICATION_PATH,
	OPT_PARAMETER,
	OPT_PARAMETER_NAME,
	OPT_PARAMETER_VALUE,
	OPT_OUTPATH,
	OPT_OUTPATH_PATH,
};

struct create_ctx {
	struct ec3_image_builder *ctx_builder;
	struct substitution_list ctx_sub_list;
	b_dict *ctx_specification;
};

static enum ec3_status create_ctx_init(
	struct create_ctx *ctx,
	const b_arglist *args)
{
	enum ec3_status status = substitution_list_init(&ctx->ctx_sub_list);

	b_arglist_option_iterator it;
	b_arglist_option_foreach_filtered(&it, args, OPT_PARAMETER)
	{
		b_arglist_value *name, *value;
		b_arglist_option_get_value(
			it.opt,
			OPT_PARAMETER_NAME,
			0,
			&name);
		b_arglist_option_get_value(
			it.opt,
			OPT_PARAMETER_VALUE,
			0,
			&value);

		status = substitution_list_add_substitution(
			&ctx->ctx_sub_list,
			name->val_str,
			value->val_str);

		if (status != EC3_SUCCESS) {
			return status;
		}
	}

	return EC3_SUCCESS;
}

static void create_ctx_finish(struct create_ctx *ctx)
{
	if (ctx->ctx_builder) {
		ec3_image_builder_destroy(ctx->ctx_builder);
	}

	if (ctx->ctx_specification) {
		b_dict_release(ctx->ctx_specification);
	}

	substitution_list_finish(&ctx->ctx_sub_list);
}

static enum ec3_status configure_image_parameters(
	struct create_ctx *ctx,
	struct ec3_parameters *param)
{
	enum ec3_status status = EC3_SUCCESS;

	b_string *ident = B_STRING(b_dict_at(ctx->ctx_specification, "id"));
	if (ident) {
		if (B_TYPEID(ident) != B_OBJECT_TYPE_STRING) {
			return EC3_ERR_BAD_FORMAT;
		}

		status = substitution_list_substitute(
			&ctx->ctx_sub_list,
			ident);
		if (status != EC3_SUCCESS) {
			return status;
		}

		status = ec3_identifier_from_string(
			b_string_ptr(ident),
			&param->p_ident);

		if (status != EC3_SUCCESS) {
			return status;
		}
	}

	b_number *cluster_size_p
		= B_NUMBER(b_dict_at(ctx->ctx_specification, "clusterSize"));
	if (B_TYPEID(cluster_size_p) != B_OBJECT_TYPE_NUMBER) {
		return EC3_ERR_BAD_FORMAT;
	}

	param->p_cluster_size = ec3_cluster_size_bytes_to_id(
		b_number_get_int(cluster_size_p));
	if (param->p_cluster_size == -1) {
		return EC3_ERR_NOT_SUPPORTED;
	}

	b_string *compression_p
		= B_STRING(b_dict_at(ctx->ctx_specification, "compression"));
	const char *compression = NULL;
	if (compression_p) {
		if (B_TYPEID(compression_p) != B_OBJECT_TYPE_STRING) {
			return EC3_ERR_BAD_FORMAT;
		}

		compression = b_string_ptr(compression_p);
		if (!strcmp(compression, "zstd")) {
			param->p_compression_func = EC3_COMPRESSION_ZSTD;
		} else {
			return EC3_ERR_NOT_SUPPORTED;
		}
	}

	b_string *encryption_p
		= B_STRING(b_dict_at(ctx->ctx_specification, "encryption"));
	const char *encryption = NULL;
	if (encryption_p) {
		if (B_TYPEID(encryption_p) != B_OBJECT_TYPE_STRING) {
			return EC3_ERR_BAD_FORMAT;
		}

		encryption = b_string_ptr(encryption_p);
		if (!strcmp(encryption, "aes256")) {
			param->p_compression_func = EC3_ENCRYPTION_AES256;
		} else {
			return EC3_ERR_NOT_SUPPORTED;
		}
	}

	return EC3_SUCCESS;
}

static b_status read_specification(const char *path_cstr, b_dict **out)
{
	b_path *path = b_path_create_from_cstr(path_cstr);
	b_file *file = NULL;
	b_status status = b_file_open(NULL, path, B_FILE_READ_ONLY, &file);
	b_path_release(path);

	if (!B_OK(status)) {
		return status;
	}

	b_stream *src = NULL;
	status = b_file_open_stream(file, &src);
	b_file_release(file);

	if (!B_OK(status)) {
		return status;
	}

	b_serial_ctx *ctx = NULL;
	b_serial_ctx_create(&ctx);

	b_object *data = NULL;
	status = b_serial_ctx_deserialise(
		ctx,
		B_SERIAL_FORMAT_JSON,
		src,
		&data,
		0);
	if (!B_OK(status)) {
		return status;
	}

	b_stream_close(src);

	if (!B_OBJECT_IS(data, DICT)) {
		b_release(data);
		return B_ERR_BAD_FORMAT;
	}

	*out = B_DICT(data);
	return B_SUCCESS;
}

#define ASSERT_OBJECT(object, status)                                          \
	if (!object) {                                                         \
		return status;                                                 \
	}
#define ASSERT_OBJECT_TYPE(object, type, status)                               \
	if (!B_OBJECT_IS(object, type)) {                                      \
		return status;                                                 \
	}

static enum ec3_status capture_volumes(struct create_ctx *ctx)
{
	enum ec3_status status = EC3_SUCCESS;

	b_array *volumes
		= B_ARRAY(b_dict_at(ctx->ctx_specification, "volumes"));
	if (!volumes) {
		return EC3_SUCCESS;
	}

	ASSERT_OBJECT_TYPE(volumes, ARRAY, EC3_ERR_BAD_FORMAT);

	b_array_iterator it;
	b_array_foreach(&it, volumes)
	{
		b_dict *volume = B_DICT(it.value);
		ASSERT_OBJECT_TYPE(volume, DICT, EC3_ERR_BAD_FORMAT);

		uint64_t id = 0;
		b_object *id_object = b_dict_at(volume, "id");
		ASSERT_OBJECT(id_object, EC3_ERR_BAD_FORMAT);

		if (B_OBJECT_IS(id_object, NUMBER)) {
			id = (uint64_t)b_number_get_int64(B_NUMBER(id_object));
		} else if (B_OBJECT_IS(id_object, STRING)) {
			substitution_list_substitute(
				&ctx->ctx_sub_list,
				B_STRING(id_object));
			status = ec3_identifier_from_string(
				b_string_ptr(B_STRING(id_object)),
				&id);
		}

		if (status != EC3_SUCCESS) {
			return status;
		}

		b_string *src_path = B_STRING(b_dict_at(volume, "sourcePath"));
		ASSERT_OBJECT(src_path, EC3_ERR_BAD_FORMAT);
		ASSERT_OBJECT_TYPE(src_path, STRING, EC3_ERR_BAD_FORMAT);
		substitution_list_substitute(&ctx->ctx_sub_list, src_path);

		char id_str[32];
		ec3_identifier_to_string(id, id_str, sizeof id_str);

		printf("adding volume 0x%" PRIx64 "(%s) -> %s\n",
		       id,
		       id_str,
		       b_string_ptr(src_path));

		status = ec3_image_builder_capture_directory(
			ctx->ctx_builder,
			id,
			b_string_ptr(src_path));
		if (status != EC3_SUCCESS) {
			b_err("cannot capture directory '%s'",
			      b_string_ptr(src_path));
			b_i("reason: %s", ec3_status_to_string(status));
			return status;
		}
	}

	return EC3_SUCCESS;
}

static enum ec3_status capture_blobs(struct create_ctx *ctx)
{
	enum ec3_status status = EC3_SUCCESS;

	b_array *blobs = B_ARRAY(b_dict_at(ctx->ctx_specification, "blobs"));
	if (!blobs) {
		return EC3_SUCCESS;
	}

	ASSERT_OBJECT_TYPE(blobs, ARRAY, EC3_ERR_BAD_FORMAT);

	b_array_iterator it;
	b_array_foreach(&it, blobs)
	{
		b_dict *blob = B_DICT(it.value);
		ASSERT_OBJECT_TYPE(blob, DICT, EC3_ERR_BAD_FORMAT);

		uint64_t id = 0;
		b_object *id_object = b_dict_at(blob, "id");
		if (!id_object) {
			return EC3_ERR_BAD_FORMAT;
		}

		if (B_OBJECT_IS(id_object, NUMBER)) {
			id = (uint64_t)b_number_get_int64(B_NUMBER(id_object));
		} else if (B_OBJECT_IS(id_object, STRING)) {
			substitution_list_substitute(
				&ctx->ctx_sub_list,
				B_STRING(id_object));
			status = ec3_identifier_from_string(
				b_string_ptr(B_STRING(id_object)),
				&id);
		}

		if (status != EC3_SUCCESS) {
			return status;
		}

		b_string *src_path = B_STRING(b_dict_at(blob, "sourcePath"));
		ASSERT_OBJECT(src_path, EC3_ERR_BAD_FORMAT);
		ASSERT_OBJECT_TYPE(src_path, STRING, EC3_ERR_BAD_FORMAT);
		substitution_list_substitute(&ctx->ctx_sub_list, src_path);

		char id_str[32];
		ec3_identifier_to_string(id, id_str, sizeof id_str);

		printf("adding blob 0x%" PRIx64 "(%s) -> %s\n",
		       id,
		       id_str,
		       b_string_ptr(src_path));

		FILE *fp = fopen(b_string_ptr(src_path), "rb");
		if (!fp) {
			b_err("cannot open file '%s'", b_string_ptr(src_path));
			b_i("reason: %s", strerror(errno));
			return EC3_ERR_NO_ENTRY;
		}

		status = ec3_image_builder_add_blob_from_file(
			ctx->ctx_builder,
			id,
			fp);

		fclose(fp);

		if (status != EC3_SUCCESS) {
			b_err("cannot capture blob '%s'",
			      b_string_ptr(src_path));
			b_i("reason: %s", ec3_status_to_string(status));
			return status;
		}
	}

	return EC3_SUCCESS;
}

static enum ec3_executable_format executable_format_string_to_id(const char *s)
{
	if (!strcmp(s, "elf")) {
		return EC3_EXEC_ELF;
	}

	return EC3_EXEC_OTHER;
}

static enum ec3_status capture_executables(struct create_ctx *ctx)
{
	enum ec3_status status = EC3_SUCCESS;

	b_array *executables
		= B_ARRAY(b_dict_at(ctx->ctx_specification, "executables"));
	if (!executables) {
		return EC3_SUCCESS;
	}

	ASSERT_OBJECT_TYPE(executables, ARRAY, EC3_ERR_BAD_FORMAT);

	b_array_iterator it;
	b_array_foreach(&it, executables)
	{
		b_dict *executable = B_DICT(it.value);
		ASSERT_OBJECT_TYPE(executable, DICT, EC3_ERR_BAD_FORMAT);

		uint64_t id = 0;
		b_object *id_object = b_dict_at(executable, "id");
		if (!id_object) {
			return EC3_ERR_BAD_FORMAT;
		}

		if (B_OBJECT_IS(id_object, NUMBER)) {
			id = (uint64_t)b_number_get_int64(B_NUMBER(id_object));
		} else if (B_OBJECT_IS(id_object, STRING)) {
			substitution_list_substitute(
				&ctx->ctx_sub_list,
				B_STRING(id_object));
			status = ec3_identifier_from_string(
				b_string_ptr(B_STRING(id_object)),
				&id);
		}

		if (status != EC3_SUCCESS) {
			return status;
		}

		b_string *src_path
			= B_STRING(b_dict_at(executable, "sourcePath"));
		ASSERT_OBJECT(src_path, EC3_ERR_BAD_FORMAT);
		ASSERT_OBJECT_TYPE(src_path, STRING, EC3_ERR_BAD_FORMAT);
		substitution_list_substitute(&ctx->ctx_sub_list, src_path);

		b_string *format = B_STRING(b_dict_at(executable, "format"));
		ASSERT_OBJECT(format, EC3_ERR_BAD_FORMAT);
		ASSERT_OBJECT_TYPE(format, STRING, EC3_ERR_BAD_FORMAT);
		substitution_list_substitute(&ctx->ctx_sub_list, format);

		b_string_tolower(format);

		char id_str[32];
		ec3_identifier_to_string(id, id_str, sizeof id_str);

		printf("adding %s executable 0x%" PRIx64 "(%s) -> %s\n",
		       b_string_ptr(format),
		       id,
		       id_str,
		       b_string_ptr(src_path));

		FILE *fp = fopen(b_string_ptr(src_path), "rb");
		if (!fp) {
			b_err("cannot open file '%s'", b_string_ptr(src_path));
			b_i("reason: %s", strerror(errno));
			return EC3_ERR_NO_ENTRY;
		}

		enum ec3_executable_format exe_format
			= executable_format_string_to_id(b_string_ptr(format));

		struct ec3_tag_executable_info exe_info = {0};
		status = get_executable_info_from_file(
			fp,
			exe_format,
			&exe_info);
		if (status != EC3_SUCCESS) {
			return status;
		}

		fseek(fp, 0, SEEK_SET);
		status = ec3_image_builder_add_executable_from_file(
			ctx->ctx_builder,
			id,
			&exe_info,
			fp);

		fclose(fp);

		if (status != EC3_SUCCESS) {
			b_err("cannot capture executable '%s'",
			      b_string_ptr(src_path));
			b_i("reason: %s", ec3_status_to_string(status));
			return status;
		}
	}

	return EC3_SUCCESS;
}

static int create(
	const b_command *self,
	const b_arglist *opt,
	const b_array *args)
{
	const char *out_path_cstr = NULL;
	const char *specification_path_cstr = NULL;
	b_status status = b_arglist_get_string(
		opt,
		OPT_SPECIFICATION,
		OPT_SPECIFICATION_PATH,
		0,
		&specification_path_cstr);
	if (!B_OK(status)) {
		b_err("no specification file specified\n");
		return -1;
	}

	status = b_arglist_get_string(
		opt,
		OPT_OUTPATH,
		OPT_OUTPATH_PATH,
		0,
		&out_path_cstr);
	if (!B_OK(status)) {
		b_err("no output file path specified\n");
		return -1;
	}

	struct create_ctx ctx = {0};
	enum ec3_status status2 = create_ctx_init(&ctx, opt);

	status = read_specification(
		specification_path_cstr,
		&ctx.ctx_specification);
	if (!B_OK(status)) {
		b_err("cannot read specification file.");
		b_i("reason: %s", b_status_to_string(status));
		return -1;
	}

	struct ec3_parameters param = {};
	status2 = configure_image_parameters(&ctx, &param);
	if (status2 != EC3_SUCCESS) {
		return -1;
	}

	status2 = ec3_image_builder_create(
		out_path_cstr,
		&param,
		&ctx.ctx_builder);

	if (status2 != EC3_SUCCESS) {
		b_err("cannot initialise EC3 writer");
		return -1;
	}

	status2 = capture_volumes(&ctx);
	if (status2 != EC3_SUCCESS) {
		return -1;
	}

	status2 = capture_blobs(&ctx);
	if (status2 != EC3_SUCCESS) {
		return -1;
	}

	status2 = capture_executables(&ctx);
	if (status2 != EC3_SUCCESS) {
		return -1;
	}

	create_ctx_finish(&ctx);
	return 0;
}

B_COMMAND(CMD_CREATE, CMD_ROOT)
{
	B_COMMAND_NAME("create");
	B_COMMAND_SHORT_NAME('C');
	B_COMMAND_DESC("create an ec3 file");
	B_COMMAND_FLAGS(B_COMMAND_SHOW_HELP_BY_DEFAULT);
	B_COMMAND_FUNCTION(create);

	B_COMMAND_HELP_OPTION();

	B_COMMAND_OPTION(OPT_OUTPATH)
	{
		B_OPTION_SHORT_NAME('o');
		B_OPTION_LONG_NAME("out");
		B_OPTION_DESC("the path to save the new file to");

		B_OPTION_ARG(OPT_OUTPATH_PATH)
		{
			B_ARG_NAME("path");
			B_ARG_NR_VALUES(1);
		}
	}

	B_COMMAND_OPTION(OPT_SPECIFICATION)
	{
		B_OPTION_SHORT_NAME('s');
		B_OPTION_LONG_NAME("specification");
		B_OPTION_DESC(
			"the path to a specification file describing the "
			"layout and parameters of the new file");

		B_OPTION_ARG(OPT_SPECIFICATION_PATH)
		{
			B_ARG_NAME("path");
			B_ARG_NR_VALUES(1);
		}
	}

	B_COMMAND_OPTION(OPT_PARAMETER)
	{
		B_OPTION_SHORT_NAME('p');
		B_OPTION_LONG_NAME("parameter");
		B_OPTION_DESC("a parameter to use when creating the new file.");

		B_OPTION_ARG(OPT_PARAMETER_NAME)
		{
			B_ARG_NAME("name");
			B_ARG_NR_VALUES(1);
		}

		B_OPTION_ARG(OPT_PARAMETER_VALUE)
		{
			B_ARG_NAME("value");
			B_ARG_NR_VALUES(1);
		}
	}

	B_COMMAND_USAGE()
	{
		B_COMMAND_USAGE_OPT(OPT_SPECIFICATION);
		B_COMMAND_USAGE_OPT(OPT_OUTPATH);
	}

	B_COMMAND_USAGE()
	{
		B_COMMAND_USAGE_OPT(OPT_SPECIFICATION);
		B_COMMAND_USAGE_OPT(OPT_PARAMETER);
		B_COMMAND_USAGE_OPT(OPT_OUTPATH);
	}
}