#include "builder.h"

#include <blue/object/hashmap.h>
#include <mie/ctx.h>
#include <mie/ir/const.h>
#include <mie/ir/data.h>
#include <mie/ir/instr.h>
#include <mie/select/builder.h>
#include <mie/select/graph.h>
#include <mie/select/node.h>
#include <mie/target/target.h>
#include <stdlib.h>
#include <string.h>

struct mie_select_builder {
	struct mie_ctx *b_ctx;
	struct mie_select_graph *b_graph;
	const struct mie_target *b_target;
	/* map of mie_instr* to mie_select_value*, defining the graph node that
	 * corresponds to each value-producing instruction */
	b_hashmap *b_nodes;
	/* map of mie_instr* to mie_select_value*, defining the graph node that
	 * last accessed a memory location defined by a particular instruction. */
	b_hashmap *b_mem_access;
};

struct mie_select_builder *mie_select_builder_create(
	struct mie_ctx *ctx, const struct mie_target *target)
{
	struct mie_select_builder *out = malloc(sizeof *out);
	if (!out) {
		return NULL;
	}

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

	out->b_ctx = ctx;
	out->b_target = target;
	out->b_graph = mie_select_graph_create(ctx);
	if (!out->b_graph) {
		free(out);
		return NULL;
	}

	out->b_nodes = b_hashmap_create(NULL, NULL);
	if (!out->b_nodes) {
		mie_select_graph_destroy(out->b_graph);
		free(out);
		return NULL;
	}

	out->b_mem_access = b_hashmap_create(NULL, NULL);
	if (!out->b_mem_access) {
		b_hashmap_release(out->b_nodes);
		mie_select_graph_destroy(out->b_graph);
		free(out);
		return NULL;
	}

	return out;
}

void mie_select_builder_destroy(struct mie_select_builder *builder)
{
	if (builder->b_nodes) {
		b_hashmap_release(builder->b_nodes);
	}

	if (builder->b_mem_access) {
		b_hashmap_release(builder->b_mem_access);
	}

	if (builder->b_graph) {
		mie_select_graph_destroy(builder->b_graph);
	}

	free(builder);
}

struct mie_select_graph *mie_select_builder_get_graph(
	struct mie_select_builder *builder)
{
	return builder->b_graph;
}

struct mie_ctx *mie_select_builder_get_ctx(struct mie_select_builder *builder)
{
	return builder->b_ctx;
}

const struct mie_target *mie_select_builder_get_target(
	struct mie_select_builder *builder)
{
	return builder->b_target;
}

static void clear_node_map(struct mie_select_builder *builder)
{
	b_hashmap_iterator it;
	b_hashmap_iterator_begin(builder->b_nodes, &it);
	while (b_hashmap_iterator_is_valid(&it)) {
		free(it.value->value_data);
		b_hashmap_iterator_erase(&it);
	}
}

struct mie_select_graph *mie_select_builder_finish(struct mie_select_builder *builder)
{
	enum mie_status status = MIE_SUCCESS;
	struct mie_select_graph *graph = builder->b_graph;

	struct mie_select_value *root_operands[] = {
		NULL,
	};
	size_t nr_root_operands = 0;

	if (graph->g_last_chain.v_node) {
		root_operands[0] = &graph->g_last_chain;
		nr_root_operands++;
	}

	status = mie_select_graph_get_node(
		graph, mie_target_builtin(), MIE_SELECT_OP_ROOT, root_operands,
		nr_root_operands, NULL, 0, &graph->g_root);

	if (status != MIE_SUCCESS) {
		return NULL;
	}

	clear_node_map(builder);

	builder->b_graph = NULL;

	return graph;
}

enum mie_status mie_select_builder_get_const(
	struct mie_select_builder *builder, long long value,
	struct mie_select_value *out)
{
	const struct mie_target *builtin = mie_target_builtin();
	struct mie_type *ctype = mie_ctx_get_int_type(builder->b_ctx, 32);
	struct mie_select_node *node = NULL;

	b_queue_iterator it;
	b_queue_foreach (&it, &builder->b_graph->g_nodes) {
		node = b_unbox(struct mie_select_node, it.entry, n_entry);

		if (node->n_target != builtin) {
			continue;
		}

		if (node->n_opcode != MIE_SELECT_OP_CONSTANT) {
			continue;
		}

		if (!(node->n_flags & MIE_SELECT_NODE_F_IVALUE)) {
			continue;
		}

		if (node->n_value.i != value) {
			continue;
		}

		mie_select_node_get_value(node, ctype, 0, out);
		return MIE_SUCCESS;
	}

	enum mie_status status = mie_select_graph_get_node(
		builder->b_graph, builtin, MIE_SELECT_OP_CONSTANT, NULL, 0,
		&ctype, 1, &node);

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

	node->n_flags |= MIE_SELECT_NODE_F_IVALUE;
	node->n_value.i = value;

	mie_select_node_get_value(node, ctype, 0, out);
	return MIE_SUCCESS;
}

enum mie_status mie_select_builder_push_instr(
	struct mie_select_builder *builder, struct mie_instr *instr)
{
	const struct select_instr_type *i_type
		= select_type_for_instr(instr->i_type);
	if (!i_type) {
		return MIE_ERR_INVALID_VALUE;
	}

	return i_type->i_push(builder, instr);
}

static enum mie_status get_const_node(
	struct mie_select_builder *builder, struct mie_value *ir_val,
	struct mie_select_value *out)
{
	struct mie_const *c = (struct mie_const *)ir_val;
	struct mie_select_node *node;
	enum mie_status status = mie_select_graph_get_node(
		builder->b_graph, mie_target_builtin(), MIE_SELECT_OP_CONSTANT,
		NULL, 0, &c->c_type, 1, &node);
	if (status != MIE_SUCCESS) {
		return status;
	}

	node->n_flags = MIE_SELECT_NODE_F_PVALUE;
	node->n_value.v = ir_val;

	mie_select_node_get_value(node, c->c_type, 0, out);

	return MIE_SUCCESS;
}

static enum mie_status get_data_node(
	struct mie_select_builder *builder, struct mie_value *ir_val,
	struct mie_select_value *out)
{
	struct mie_data *data = (struct mie_data *)ir_val;
	unsigned int opcode = 0;
	struct mie_type *type = NULL;

	switch (data->d_type) {
	case MIE_DATA_EXTERN_GLOBAL:
		opcode = MIE_SELECT_OP_GLOBAL_ADDRESS;
		type = mie_ctx_get_type(builder->b_ctx, MIE_TYPE_PTR);
		break;
	default:
		return MIE_ERR_INVALID_VALUE;
	}

	struct mie_select_node *node;
	enum mie_status status = mie_select_graph_get_node(
		builder->b_graph, mie_target_builtin(), opcode, NULL, 0, &type,
		1, &node);
	if (status != MIE_SUCCESS) {
		return status;
	}

	node->n_flags = MIE_SELECT_NODE_F_PVALUE;
	node->n_value.v = ir_val;

	mie_select_node_get_value(node, type, 0, out);

	return MIE_SUCCESS;
}

struct mie_select_value *mie_select_builder_get_value(
	struct mie_select_builder *builder, struct mie_value *ir_val)
{
	b_hashmap_key key = {
		.key_flags = B_HASHMAP_KEY_F_INTVALUE,
		.key_data = ir_val,
		.key_size = sizeof(struct mie_value *),
	};

	const b_hashmap_value *val = b_hashmap_get(builder->b_nodes, &key);
	if (val) {
		return val->value_data;
	}

	struct mie_select_value *select_val = malloc(sizeof *select_val);
	if (!select_val) {
		return NULL;
	}

	enum mie_status status = MIE_ERR_INVALID_VALUE;

	switch (ir_val->v_type->t_id) {
	case MIE_VALUE_CONST:
		status = get_const_node(builder, ir_val, select_val);
		break;
	case MIE_VALUE_DATA:
		status = get_data_node(builder, ir_val, select_val);
		break;
	default:
		break;
	}

	if (status != MIE_SUCCESS) {
		return NULL;
	}

	key.key_data = ir_val;
	key.key_size = sizeof(struct mie_value *);

	b_hashmap_value hashmap_val = {
		.value_data = select_val,
		.value_size = sizeof *select_val,
	};
	b_hashmap_put(builder->b_nodes, &key, &hashmap_val);

	return select_val;
}

enum mie_status mie_select_builder_set_value(
	struct mie_select_builder *builder, struct mie_value *ir_val,
	struct mie_select_value *graph_val)
{
	struct mie_select_value *graph_val2 = malloc(sizeof *graph_val2);
	if (!graph_val2) {
		return MIE_ERR_NO_MEMORY;
	}

	memcpy(graph_val2, graph_val, sizeof *graph_val);

	b_hashmap_key key = {
		.key_flags = B_HASHMAP_KEY_F_INTVALUE,
		.key_data = ir_val,
		.key_size = sizeof(struct mie_value *),
	};

	b_hashmap_value hashmap_val = {
		.value_data = graph_val2,
		.value_size = sizeof *graph_val2,
	};

	b_hashmap_put(builder->b_nodes, &key, &hashmap_val);
	return MIE_SUCCESS;
}

struct mie_select_value *mie_select_builder_get_mem_access(
	struct mie_select_builder *builder, struct mie_value *ir_val)
{
	b_hashmap_key key = {
		.key_flags = B_HASHMAP_KEY_F_INTVALUE,
		.key_data = ir_val,
		.key_size = sizeof(struct mie_value *),
	};

	const b_hashmap_value *val = b_hashmap_get(builder->b_mem_access, &key);
	if (val) {
		return val->value_data;
	}

	return NULL;
}

enum mie_status mie_select_builder_set_mem_access(
	struct mie_select_builder *builder, struct mie_value *ir_val,
	struct mie_select_value *graph_val)
{
	struct mie_select_value *graph_val2 = malloc(sizeof *graph_val2);
	if (!graph_val2) {
		return MIE_ERR_NO_MEMORY;
	}

	memcpy(graph_val2, graph_val, sizeof *graph_val);

	b_hashmap_key key = {
		.key_flags = B_HASHMAP_KEY_F_INTVALUE,
		.key_data = ir_val,
		.key_size = sizeof(struct mie_value *),
	};

	b_hashmap_value hashmap_val = {
		.value_data = graph_val2,
		.value_size = sizeof *graph_val2,
	};

	b_hashmap_put(builder->b_mem_access, &key, &hashmap_val);
	return MIE_SUCCESS;
}

struct mie_select_node *mie_select_builder_find_node_with_ivalue(
	struct mie_select_builder *builder, const struct mie_target *target,
	unsigned int opcode, long long val)
{
	b_queue_iterator it;
	b_queue_foreach (&it, &builder->b_graph->g_nodes) {
		struct mie_select_node *node
			= b_unbox(struct mie_select_node, it.entry, n_entry);

		if (node->n_target == target && node->n_opcode == opcode
		    && node->n_value.i == val) {
			return node;
		}
	}

	return NULL;
}