#include "builder.h"

#include "fs-tree.h"
#include "image.h"
#include "tag.h"
#include "volume.h"

#include <blue/io/directory.h>
#include <blue/io/path.h>
#include <blue/object/list.h>
#include <stdlib.h>
#include <string.h>

struct capture_directory_ctx {
	struct ec3_image_builder *ctx_builder;
	struct ec3_volume *ctx_volume;
	struct fs_tree ctx_tree;
};

static enum ec3_status init_string_table(struct ec3_image_builder *builder)
{
	if (builder->b_flags & EC3_IMAGE_BUILDER_STRING_TABLE) {
		return EC3_SUCCESS;
	}

	uint64_t stab_id;
	ec3_identifier_from_string("_VOLSTR0", &stab_id);

	enum ec3_status status = ec3_image_ioctx_create_tag(
		builder->b_image,
		EC3_TAG_STAB,
		stab_id,
		EC3_TAG_IO_WRITE | EC3_TAG_IO_SEQUENTIAL,
		&builder->b_stab);
	if (status != EC3_SUCCESS) {
		return status;
	}

	string_table_init(&builder->b_strings);
	builder->b_flags |= EC3_IMAGE_BUILDER_STRING_TABLE;

	return EC3_SUCCESS;
}

static enum ec3_status init_chunk_table(struct ec3_image_builder *builder)
{
	if (builder->b_flags & EC3_IMAGE_BUILDER_CHUNK_TABLE) {
		return EC3_SUCCESS;
	}

	uint64_t ctab_id, cdat_id;
	ec3_identifier_from_string("_CHKTAB0", &ctab_id);
	ec3_identifier_from_string("_CHKDAT0", &cdat_id);

	enum ec3_status status = ec3_image_ioctx_create_tag(
		builder->b_image,
		EC3_TAG_CTAB,
		ctab_id,
		EC3_TAG_IO_READWRITE,
		&builder->b_ctab);
	if (status != EC3_SUCCESS) {
		return status;
	}

	status = ec3_image_ioctx_create_tag(
		builder->b_image,
		EC3_TAG_CDAT,
		cdat_id,
		EC3_TAG_IO_READWRITE,
		&builder->b_cdat);
	if (status != EC3_SUCCESS) {
		ec3_tag_ioctx_close(builder->b_ctab);
		builder->b_ctab = NULL;
		return status;
	}

	status = chunk_table_init(
		builder->b_ctab,
		builder->b_cdat,
		builder->b_param.p_cluster_size,
		&builder->b_chunks);
	if (status != EC3_SUCCESS) {
		ec3_tag_ioctx_close(builder->b_ctab);
		ec3_tag_ioctx_close(builder->b_cdat);
		return status;
	}

	chunk_table_init_empty_table(&builder->b_chunks);

	builder->b_flags |= EC3_IMAGE_BUILDER_CHUNK_TABLE;

	return EC3_SUCCESS;
}

static enum ec3_status capture_directory_file(
	struct capture_directory_ctx *ctx,
	b_directory *dir,
	const char *filename,
	const b_path *filepath,
	ec3_chunk_id out_chunk)
{
	const struct ec3_image_info *image_info
		= ec3_image_ioctx_get_info(ctx->ctx_builder->b_image);
	size_t key = string_table_get(&ctx->ctx_builder->b_strings, filename);
	size_t buf_len = image_info->img_cluster_size;
	char *buf = malloc(buf_len);

	if (!buf) {
		return EC3_ERR_NO_MEMORY;
	}

	b_file *src = NULL;
	b_status status = b_file_open(
		dir,
		filepath,
		B_FILE_READ_ONLY | B_FILE_BINARY,
		&src);
	if (!B_OK(status)) {
		free(buf);
		return ec3_status_from_b_status(status, EC3_ERR_NO_ENTRY);
	}

	enum ec3_status s2 = EC3_SUCCESS;
	chunk_table_begin_chunk(&ctx->ctx_builder->b_chunks);

	size_t chunk_size = 0;

	while (1) {
		size_t nr_read = 0;
		status = b_file_read(
			src,
			B_OFFSET_CURRENT,
			buf_len,
			buf,
			&nr_read);

		if (!B_OK(status)) {
			s2 = ec3_status_from_b_status(
				status,
				EC3_ERR_IO_FAILURE);
			break;
		}

		enum ec3_status s2 = chunk_table_put(
			&ctx->ctx_builder->b_chunks,
			buf,
			nr_read);

		if (s2 != EC3_SUCCESS) {
			break;
		}

		chunk_size += nr_read;

		if (nr_read < buf_len) {
			break;
		}
	}

	s2 = chunk_table_end_chunk(&ctx->ctx_builder->b_chunks, out_chunk);

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

	char id_str[128];
	ec3_chunk_id_to_string(out_chunk, id_str, sizeof id_str);

	free(buf);
	b_file_release(src);
	return s2;
}

struct capture_directory_structure_args {
	b_list *args_stack;
	struct capture_directory_ctx *args_ctx;
};

static int capture_directory_structure_callback(
	struct fs_tree *tree,
	struct fs_tree_node *node,
	unsigned int depth,
	int direction,
	void *arg)
{
	if (direction != ITERATE_POSTORDER) {
		return 0;
	}

	struct capture_directory_structure_args *args = arg;
	b_list *stack = args->args_stack;
	struct capture_directory_ctx *ctx = args->args_ctx;

	if (node->n_type == FS_TREE_FILE) {
		b_list_push_back(stack, node);
		return 0;
	}

	size_t nr_children = fs_tree_node_get_nr_children(node);
	struct ec3_directory_entry dent;
	enum ec3_status status = EC3_SUCCESS;

	chunk_table_begin_chunk(&ctx->ctx_builder->b_chunks);
	for (size_t i = 0; i < nr_children; i++) {
		struct fs_tree_node *child = b_list_pop_back(stack);
		if (!child) {
			return EC3_ERR_INTERNAL_FAILURE;
		}

		dent.d_name = b_i32_htob(string_table_get(
			&ctx->ctx_builder->b_strings,
			child->n_name));
		dent.d_vnode = b_i32_htob(child->n_id);

		status = chunk_table_put(
			&ctx->ctx_builder->b_chunks,
			&dent,
			sizeof dent);
	}

	status = chunk_table_end_chunk(
		&ctx->ctx_builder->b_chunks,
		node->n_chunk);
	if (status != EC3_SUCCESS) {
		return status;
	}

	char id_str[128];
	ec3_chunk_id_to_string(node->n_chunk, id_str, sizeof id_str);
#if 0
	printf("[d %s / %zu] wrote %zu byte chunk %s\n",
	       node->n_name,
	       node->n_id,
	       nr_children * sizeof dent,
	       id_str);
#endif

	struct ec3_vnode vnode = {
		.v_id = node->n_id,
		.v_mode = EC3_V_DIR,
	};

	memcpy(vnode.v_data, node->n_chunk, sizeof vnode.v_data);

	ec3_volume_put_vnode(ctx->ctx_volume, &vnode);

	b_list_push_back(stack, node);

	return 0;
}

static enum ec3_status capture_directory_structure(
	struct capture_directory_ctx *ctx)
{
	b_list *stack = b_list_create();

	struct capture_directory_structure_args args = {
		.args_stack = stack,
		.args_ctx = ctx,
	};

	fs_tree_iterate(
		&ctx->ctx_tree,
		capture_directory_structure_callback,
		&args);

	return EC3_SUCCESS;
}

enum ec3_status ec3_image_builder_capture_directory(
	struct ec3_image_builder *builder,
	uint64_t id,
	const char *directory_path)
{
	enum ec3_status status2 = init_string_table(builder);
	if (status2 != EC3_SUCCESS) {
		return status2;
	}

	status2 = init_chunk_table(builder);
	if (status2 != EC3_SUCCESS) {
		return status2;
	}

	b_path *path = b_path_create_from_cstr(directory_path);
	b_directory *dir;
	b_status status = b_directory_open(NULL, path, &dir);
	b_path_release(path);

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

	struct ec3_tag_ioctx *volu;
	enum ec3_status s2 = ec3_image_ioctx_create_tag(
		builder->b_image,
		EC3_TAG_VOLU,
		id,
		EC3_TAG_IO_READ | EC3_TAG_IO_WRITE,
		&volu);
	if (s2 != EC3_SUCCESS) {
		b_directory_release(dir);
		return s2;
	}

	struct capture_directory_ctx ctx = {.ctx_builder = builder};

	s2 = ec3_volume_create(builder->b_image, volu, &ctx.ctx_volume);

	if (s2 != EC3_SUCCESS) {
		ec3_tag_ioctx_close(volu);
		b_directory_release(dir);
		return s2;
	}

	fs_tree_init(&ctx.ctx_tree);

	struct ec3_vnode vnode = {0};
	vnode.v_id = ctx.ctx_tree.fs_root->n_id;
	ec3_volume_put_vnode(ctx.ctx_volume, &vnode);

	status2 = EC3_SUCCESS;
	b_directory_iterator it;
	b_directory_iterator_begin(dir, &it, B_DIRECTORY_ITERATE_PARENT_LAST);
	while (b_directory_iterator_is_valid(&it)) {
		if (!b_directory_path_is_file(dir, it.filepath)) {
			printf("dir: %s\n", b_path_ptr(it.filepath));
			b_directory_iterator_next(&it);
			continue;
		}

		ec3_vnode_from_file_info(&it.info, &vnode);

		printf("file: %s\n", b_path_ptr(it.filepath));

		status2 = capture_directory_file(
			&ctx,
			dir,
			it.filename,
			it.filepath,
			vnode.v_data);

		if (status2 != EC3_SUCCESS) {
#if 0
			b_err("failed to capture file '%s'",
			      b_path_ptr(it.filepath));
			b_i("error code: %s", ec3_status_to_string(status2));
#endif
			break;
		}

		fs_tree_put(
			&ctx.ctx_tree,
			b_path_ptr(it.filepath),
			&vnode.v_id,
			vnode.v_data);

		ec3_volume_put_vnode(ctx.ctx_volume, &vnode);

		b_directory_iterator_next(&it);
	}

	// fs_tree_iterate(&ctx->ctx_tree, print_fs_tree_node, NULL);

	capture_directory_structure(&ctx);
	fs_tree_fini(&ctx.ctx_tree);

	ec3_tag_ioctx_close(volu);

	b_directory_release(dir);
	return status2;
}

enum ec3_status ec3_image_builder_create(
	const char *path,
	const struct ec3_parameters *params,
	struct ec3_image_builder **out)
{
	struct ec3_image_builder *builder = malloc(sizeof *builder);
	if (!builder) {
		return EC3_ERR_NO_MEMORY;
	}

	memset(builder, 0x0, sizeof *builder);

	b_path *image_path = b_path_create_from_cstr(path);
	if (b_path_exists(image_path)) {
		b_path_unlink(image_path);
	}
	b_path_release(image_path);

	enum ec3_status status = ec3_image_ioctx_open(
		path,
		params,
		EC3_IMAGE_IO_WRITE,
		&builder->b_image);
	if (status != EC3_SUCCESS) {
		ec3_image_builder_destroy(builder);
		return status;
	}

	memcpy(&builder->b_param, params, sizeof *params);
	*out = builder;

	return EC3_SUCCESS;
}

static enum ec3_status flush_strings(struct ec3_image_builder *builder)
{
	enum ec3_status status = EC3_SUCCESS;

	b_btree_iterator s_it;
	b_btree_iterator_begin(&builder->b_strings.s_offset_tree, &s_it);
	while (b_btree_iterator_is_valid(&s_it)) {
		struct string_table_entry *entry = b_unbox(
			struct string_table_entry,
			s_it.node,
			e_offset_node);

		size_t len = strlen(entry->e_str) + 1;
		size_t nr_written = 0;
		status = ec3_tag_ioctx_write(
			builder->b_stab,
			entry->e_str,
			len,
			&nr_written);
		b_btree_iterator_next(&s_it);

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

	return status;
}

enum ec3_status ec3_image_builder_destroy(struct ec3_image_builder *builder)
{
	enum ec3_status status = EC3_SUCCESS;
	if (builder->b_flags & EC3_IMAGE_BUILDER_STRING_TABLE) {
		status = flush_strings(builder);
		string_table_finish(&builder->b_strings);
		ec3_tag_ioctx_close(builder->b_stab);
	}

	/* TODO propagate status code from flush_strings */

	if (builder->b_flags & EC3_IMAGE_BUILDER_CHUNK_TABLE) {
		chunk_table_finish(&builder->b_chunks);
		ec3_tag_ioctx_close(builder->b_ctab);
		ec3_tag_ioctx_close(builder->b_cdat);
	}

	status = ec3_image_ioctx_close(builder->b_image);
	free(builder);

	return status;
}

enum ec3_status ec3_image_builder_add_blob_from_file(
	struct ec3_image_builder *builder,
	uint64_t id,
	FILE *fp)
{
	size_t cluster_size
		= ec3_cluster_size_id_to_bytes(builder->b_param.p_cluster_size);
	char *buf = malloc(cluster_size);
	if (!buf) {
		return EC3_ERR_NO_MEMORY;
	}

	struct ec3_tag_ioctx *tag = NULL;
	enum ec3_status status = ec3_image_ioctx_create_tag(
		builder->b_image,
		EC3_TAG_BLOB,
		id,
		EC3_TAG_IO_WRITE | EC3_TAG_IO_SEQUENTIAL,
		&tag);

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

	size_t i = 0;

	while (1) {
		size_t r = fread(buf, 1, cluster_size, fp);
		if (r == 0) {
			break;
		}

		size_t w;
		status = ec3_tag_ioctx_write_cluster(tag, i++, buf, r, &w);

		if (r < cluster_size) {
			break;
		}
	}

	free(buf);
	ec3_tag_ioctx_close(tag);

	return EC3_SUCCESS;
}

enum ec3_status ec3_image_builder_add_blob_from_buffer(
	struct ec3_image_builder *builder,
	uint64_t id,
	const void *buf,
	size_t len)
{
	size_t cluster_size
		= ec3_cluster_size_id_to_bytes(builder->b_param.p_cluster_size);

	struct ec3_tag_ioctx *tag = NULL;
	enum ec3_status status = ec3_image_ioctx_create_tag(
		builder->b_image,
		EC3_TAG_BLOB,
		id,
		EC3_TAG_IO_WRITE | EC3_TAG_IO_SEQUENTIAL,
		&tag);

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

	size_t i = 0;
	const unsigned char *bytes = buf;
	size_t remaining = len;

	while (remaining > 0) {
		size_t to_copy = remaining;
		if (to_copy > cluster_size) {
			to_copy = cluster_size;
		}

		size_t w;
		status = ec3_tag_ioctx_write_cluster(
			tag,
			i++,
			bytes,
			to_copy,
			&w);

		remaining -= to_copy;
		bytes += to_copy;
	}

	ec3_tag_ioctx_close(tag);

	return EC3_SUCCESS;
}

enum ec3_status ec3_image_builder_add_executable_from_file(
	struct ec3_image_builder *builder,
	uint64_t id,
	const struct ec3_tag_executable_info *exe,
	FILE *fp)
{
	size_t cluster_size
		= ec3_cluster_size_id_to_bytes(builder->b_param.p_cluster_size);
	char *buf = malloc(cluster_size);
	if (!buf) {
		return EC3_ERR_NO_MEMORY;
	}

	struct ec3_tag_ioctx *tag = NULL;
	enum ec3_status status = ec3_image_ioctx_create_tag(
		builder->b_image,
		EC3_TAG_EXEC,
		id,
		EC3_TAG_IO_WRITE | EC3_TAG_IO_SEQUENTIAL,
		&tag);

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

	size_t i = 0;

	while (1) {
		size_t r = fread(buf, 1, cluster_size, fp);
		if (r == 0) {
			break;
		}

		size_t w;
		status = ec3_tag_ioctx_write_cluster(tag, i++, buf, r, &w);

		if (r < cluster_size) {
			break;
		}
	}

	memcpy(&tag->io_tag_info.tag_exe, exe, sizeof *exe);

	free(buf);
	ec3_tag_ioctx_close(tag);

	return EC3_SUCCESS;
}