#include <blue/core/btree.h>
#include <mie/ctx.h>
#include <mie/dialect/arith.h>
#include <stdlib.h>
#include <string.h>

typedef union float_key {
	double f;
	unsigned char i[sizeof(double)];
} float_key;

struct float_width_cache_entry {
	b_btree_node e_node;
	float_key e_key;
	struct mie_float e_value;
};

struct float_width_cache {
	b_btree_node c_node;
	size_t c_width;
	b_btree c_floats;
};

struct mie_float_cache {
	b_btree c_widths;
};

static int float_key_compare(const float_key *a, const float_key *b)
{
	return memcmp(a->i, b->i, sizeof a->i);
}

static B_BTREE_DEFINE_SIMPLE_INSERT(
	struct float_width_cache, c_node, c_width, put_width_cache);
static B_BTREE_DEFINE_SIMPLE_GET(
	struct float_width_cache, size_t, c_node, c_width, get_width_cache);

void put_float(b_btree *tree, struct float_width_cache_entry *node)
{
	if (!tree->b_root) {
		tree->b_root = &node->e_node;
		b_btree_insert_fixup(tree, &node->e_node);
		return;
	}

	b_btree_node *cur = tree->b_root;
	while (1) {
		struct float_width_cache_entry *cur_node
			= b_unbox(struct float_width_cache_entry, cur, e_node);
		b_btree_node *next = NULL;
		int cmp = float_key_compare(&node->e_key, &cur_node->e_key);

		if (cmp >= 0) {
			next = b_btree_right(cur);

			if (!next) {
				b_btree_put_right(cur, &node->e_node);
				break;
			}
		} else if (cmp < 0) {
			next = b_btree_left(cur);

			if (!next) {
				b_btree_put_left(cur, &node->e_node);
				break;
			}
		} else {
			return;
		}

		cur = next;
	}

	b_btree_insert_fixup(tree, &node->e_node);
}

static struct float_width_cache_entry *get_float(const b_btree *tree, float_key key)
{
	b_btree_node *cur = tree->b_root;
	while (cur) {
		struct float_width_cache_entry *cur_node
			= b_unbox(struct float_width_cache_entry, cur, e_node);
		int cmp = float_key_compare(&key, &cur_node->e_key);
		if (cmp > 0) {
			cur = b_btree_right(cur);
		} else if (cmp < 0) {
			cur = b_btree_left(cur);
		} else {
			return cur_node;
		}
	}

	return NULL;
}

static struct float_width_cache *float_width_cache_create(size_t width)
{
	struct float_width_cache *out = malloc(sizeof *out);
	if (!out) {
		return NULL;
	}

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

	out->c_width = width;

	return out;
}

static void float_width_cache_destroy(struct float_width_cache *cache)
{
	b_btree_node *cur = b_btree_first(&cache->c_floats);
	while (cur) {
		b_btree_node *next = b_btree_next(cur);
		b_btree_delete(&cache->c_floats, cur);

		struct float_width_cache_entry *entry
			= b_unbox(struct float_width_cache_entry, cur, e_node);
		free(entry);

		cur = next;
	}

	free(cache);
}

static struct float_width_cache_entry *float_width_cache_entry_create(
	struct mie_ctx *ctx, size_t width, double value)
{
	struct float_width_cache_entry *out = malloc(sizeof *out);
	if (!out) {
		return NULL;
	}

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

	switch (width) {
	case MIE_FLOAT_32:
		out->e_value.f_val.v_32 = value;
		break;
	case MIE_FLOAT_64:
		out->e_value.f_val.v_64 = value;
		break;
	default:
		free(out);
		return NULL;
	}

	out->e_value.f_base.a_def
		= mie_ctx_get_attribute_definition(ctx, "arith", "float");
	out->e_value.f_type = mie_arith_float_get_type(ctx, width);
	if (!out->e_value.f_type) {
		free(out);
		return NULL;
	}

	return out;
}

struct mie_float_cache *mie_float_cache_create(void)
{
	struct mie_float_cache *out = malloc(sizeof *out);
	if (!out) {
		return NULL;
	}

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

void mie_float_cache_destroy(struct mie_float_cache *cache)
{
	b_btree_node *cur = b_btree_first(&cache->c_widths);
	while (cur) {
		b_btree_node *next = b_btree_next(cur);
		b_btree_delete(&cache->c_widths, cur);

		struct float_width_cache *width_cache
			= b_unbox(struct float_width_cache, cur, c_node);
		float_width_cache_destroy(width_cache);

		cur = next;
	}

	free(cache);
}

struct mie_float *mie_float_cache_get(
	struct mie_float_cache *cache, struct mie_ctx *ctx, double value,
	size_t bit_width)
{
	struct float_width_cache *width_cache
		= get_width_cache(&cache->c_widths, bit_width);
	if (!width_cache) {
		width_cache = float_width_cache_create(bit_width);

		if (!width_cache) {
			return NULL;
		}

		put_width_cache(&cache->c_widths, width_cache);
	}

	float_key key = {.f = value};
	struct float_width_cache_entry *entry
		= get_float(&width_cache->c_floats, key);
	if (!entry) {
		entry = float_width_cache_entry_create(ctx, bit_width, value);

		if (!entry) {
			return NULL;
		}

		put_float(&width_cache->c_floats, entry);
	}

	return &entry->e_value;
}