#include <blue/object/hashmap.h>
#include <blue/object/list.h>
#include <blue/object/string.h>
#include <mie/ctx.h>
#include <mie/ir/const.h>
#include <stdlib.h>
#include <string.h>

struct ctx_int_cache_entry {
	b_btree_node i_node;
	struct mie_type i_type;
	b_btree i_values;
};

struct ctx_int_value_cache_entry {
	b_btree_node i_node;
	struct mie_int i_value;
};

B_BTREE_DEFINE_SIMPLE_INSERT(
	struct ctx_int_cache_entry, i_node, i_type.t_width, put_cached_int_type)
B_BTREE_DEFINE_SIMPLE_GET(
	struct ctx_int_cache_entry, unsigned int, i_node, i_type.t_width,
	get_cached_int_type)

B_BTREE_DEFINE_SIMPLE_INSERT(
	struct ctx_int_value_cache_entry, i_node, i_value.i_value,
	put_cached_int_value)
B_BTREE_DEFINE_SIMPLE_GET(
	struct ctx_int_value_cache_entry, int64_t, i_node, i_value.i_value,
	get_cached_int_value)

struct mie_ctx *mie_ctx_create(void)
{
	struct mie_ctx *out = malloc(sizeof *out);

	if (!out) {
		return NULL;
	}

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

	out->ctx_true = MIE_CONST(mie_ctx_get_int(out, 1, 1));
	out->ctx_false = MIE_CONST(mie_ctx_get_int(out, 0, 1));

	out->ctx_sel_cache = b_hashmap_create(free, free);
	out->ctx_string_cache = b_hashmap_create(free, free);

	return out;
}

void mie_ctx_destroy(struct mie_ctx *ctx)
{
	ctx->ctx_true = NULL;
	ctx->ctx_false = NULL;

	b_btree_iterator it = {};
	b_btree_iterator_begin(&ctx->ctx_int_cache, &it);
	while (b_btree_iterator_is_valid(&it)) {
		struct ctx_int_cache_entry *entry
			= b_unbox(struct ctx_int_cache_entry, it.node, i_node);
		b_btree_iterator_erase(&it);

		b_btree_iterator it2 = {};
		b_btree_iterator_begin(&entry->i_values, &it2);
		while (b_btree_iterator_is_valid(&it2)) {
			struct ctx_int_value_cache_entry *value = b_unbox(
				struct ctx_int_value_cache_entry, it2.node, i_node);
			b_btree_iterator_erase(&it2);
			free(value);
		}

		free(entry);
	}

	const size_t nr_types = sizeof ctx->ctx_types / sizeof ctx->ctx_types[0];
	for (size_t i = 0; i < nr_types; i++) {
		if (ctx->ctx_types[i]) {
			mie_value_destroy(MIE_VALUE(ctx->ctx_types[i]));
			ctx->ctx_types[i] = NULL;
		}
	}

	b_hashmap_release(ctx->ctx_sel_cache);

	free(ctx);
}

struct mie_type *mie_ctx_get_type(struct mie_ctx *ctx, enum mie_type_id type_id)
{
	if (type_id == MIE_TYPE_INT) {
		return NULL;
	}

	if (ctx->ctx_types[type_id]) {
		return ctx->ctx_types[type_id];
	}

	struct mie_type *type = mie_type_create();
	if (!type) {
		return NULL;
	}

	type->t_id = type_id;

	ctx->ctx_types[type_id] = type;
	return type;
}

struct mie_type *mie_ctx_get_int_type(struct mie_ctx *ctx, unsigned int nr_bits)
{
	struct ctx_int_cache_entry *entry
		= get_cached_int_type(&ctx->ctx_int_cache, nr_bits);

	if (entry) {
		return &entry->i_type;
	}

	entry = malloc(sizeof *entry);
	if (!entry) {
		return NULL;
	}

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

	entry->i_type.t_id = MIE_TYPE_INT;
	entry->i_type.t_width = nr_bits;

	put_cached_int_type(&ctx->ctx_int_cache, entry);
	return &entry->i_type;
}

struct mie_value *mie_ctx_get_bool(struct mie_ctx *ctx, bool val)
{
	return MIE_VALUE(val ? ctx->ctx_true : ctx->ctx_false);
}

struct mie_value *mie_ctx_get_int(
	struct mie_ctx *ctx, long long val, unsigned int nr_bits)
{
	struct ctx_int_cache_entry *entry
		= get_cached_int_type(&ctx->ctx_int_cache, nr_bits);

	if (!entry) {
		entry = malloc(sizeof *entry);
		if (!entry) {
			return NULL;
		}

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

		entry->i_type.t_id = MIE_TYPE_INT;
		entry->i_type.t_width = nr_bits;

		put_cached_int_type(&ctx->ctx_int_cache, entry);
	}

	struct ctx_int_value_cache_entry *value
		= get_cached_int_value(&entry->i_values, val);
	if (value) {
		return MIE_VALUE(&value->i_value);
	}

	value = malloc(sizeof *value);
	if (!value) {
		return NULL;
	}

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

	mie_const_init(&value->i_value.i_base, &entry->i_type);
	value->i_value.i_value = val;

	put_cached_int_value(&entry->i_values, value);

	return MIE_VALUE(&value->i_value);
}

struct mie_value *mie_ctx_get_selector(struct mie_ctx *ctx, const char *sel)
{
	b_hashmap_key key = {
		.key_data = sel,
		.key_size = strlen(sel),
	};

	const b_hashmap_value *cache_entry
		= b_hashmap_get(ctx->ctx_sel_cache, &key);

	if (cache_entry) {
		return cache_entry->value_data;
	}

	struct mie_selector *sel_value = malloc(sizeof *sel_value);
	if (!sel_value) {
		return NULL;
	}

	struct mie_type *sel_type = mie_ctx_get_type(ctx, MIE_TYPE_SELECTOR);
	mie_const_init(&sel_value->sel_base, sel_type);
	sel_value->sel_value = b_strdup(sel);

	key.key_data = sel_value->sel_value;

	b_hashmap_value hashmap_value = {
		.value_data = sel_value,
		.value_size = sizeof *sel_value,
	};

	b_hashmap_put(ctx->ctx_sel_cache, &key, &hashmap_value);

	return MIE_VALUE(sel_value);
}

struct mie_value *mie_ctx_get_string(struct mie_ctx *ctx, const char *s)
{
	b_hashmap_key key = {
		.key_data = s,
		.key_size = strlen(s),
	};

	const b_hashmap_value *cache_entry
		= b_hashmap_get(ctx->ctx_string_cache, &key);

	if (cache_entry) {
		return cache_entry->value_data;
	}

	struct mie_string *string_value = malloc(sizeof *string_value);
	if (!string_value) {
		return NULL;
	}

	struct mie_type *string_type = mie_ctx_get_type(ctx, MIE_TYPE_STR);
	mie_const_init(&string_value->s_base, string_type);
	string_value->s_value = b_strdup(s);

	key.key_data = string_value->s_value;

	b_hashmap_value hashmap_value = {
		.value_data = string_value,
		.value_size = sizeof *string_value,
	};

	b_hashmap_put(ctx->ctx_string_cache, &key, &hashmap_value);

	return MIE_VALUE(string_value);
}

struct mie_value *mie_ctx_create_array(struct mie_ctx *ctx)
{
	struct mie_type *array_type = mie_ctx_get_type(ctx, MIE_TYPE_ARRAY);
	struct mie_array *array = malloc(sizeof *array);

	if (!array) {
		return NULL;
	}

	memset(array, 0x0, sizeof *array);
	mie_const_init(&array->a_base, array_type);

	array->a_values = b_list_create();

	return MIE_VALUE(array);
}