#include "builder.h"

#include <blue/ds/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_unref(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_unref(builder->b_nodes);
	}

	if (builder->b_mem_access) {
		b_hashmap_unref(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;
}

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;

	struct mie_select_value incoming_chain = {};
	status = mie_select_builder_collapse_chain_ends(builder, &incoming_chain);
	if (status != MIE_SUCCESS) {
		return NULL;
	}

	if (incoming_chain.v_node) {
		root_operands[0] = &incoming_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;
	}

	b_hashmap_unref(builder->b_nodes);
	builder->b_nodes = b_hashmap_create(NULL, NULL);

	b_hashmap_unref(builder->b_mem_access);
	builder->b_mem_access = b_hashmap_create(NULL, NULL);

	builder->b_graph = mie_select_graph_create(builder->b_ctx);

	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_entry *entry = b_queue_first(&builder->b_graph->g_nodes);
	while (entry) {
		node = b_unbox(struct mie_select_node, entry, n_entry);

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

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

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

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

		mie_select_node_get_value(node, ctype, 0, out);
		return MIE_SUCCESS;
	skip:
		entry = b_queue_next(entry);
	}

	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;
}

static enum mie_status get_external_value_node(
	struct mie_select_builder *builder, struct mie_value *ir_val,
	struct mie_select_value *out)
{
	struct mie_type *type = mie_value_get_type(ir_val, builder->b_ctx);

	if (!type) {
		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(), MIE_SELECT_OP_REGISTER,
		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:
		status = get_external_value_node(builder, ir_val, select_val);
		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;
}

enum mie_status mie_select_builder_collapse_chain_ends(
	struct mie_select_builder *builder, struct mie_select_value *out)
{
	size_t nr_chains = b_queue_length(&builder->b_graph->g_chain_ends);
	b_queue_entry *entry = NULL;
	struct mie_select_chain_end *end = NULL;

	switch (nr_chains) {
	case 0:
		memset(out, 0x0, sizeof *out);
		return MIE_SUCCESS;
	case 1:
		entry = b_queue_first(&builder->b_graph->g_chain_ends);
		end = b_unbox(struct mie_select_chain_end, entry, c_entry);
		memcpy(out, end, sizeof *out);
		return MIE_SUCCESS;
	default:
		break;
	}

	struct mie_type *chain_type
		= mie_ctx_get_type(builder->b_ctx, MIE_TYPE_OTHER);

	struct mie_select_value **operands = calloc(nr_chains, sizeof *operands);
	if (!operands) {
		return MIE_ERR_NO_MEMORY;
	}

	entry = b_queue_first(&builder->b_graph->g_chain_ends);
	size_t i = 0;
	while (entry) {
		end = b_unbox(struct mie_select_chain_end, entry, c_entry);
		operands[i++] = &end->c_value;
		entry = b_queue_next(entry);
	}

	struct mie_select_node *group = NULL;
	enum mie_status status = mie_select_graph_get_node(
		builder->b_graph, mie_target_builtin(), MIE_SELECT_OP_CHAIN_GROUP,
		operands, nr_chains, &chain_type, 1, &group);
	free(operands);

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

	entry = b_queue_first(&builder->b_graph->g_chain_ends);
	while (entry) {
		end = b_unbox(struct mie_select_chain_end, entry, c_entry);
		b_queue_entry *next = b_queue_next(entry);

		b_queue_delete(&builder->b_graph->g_chain_ends, entry);
		free(end);

		entry = next;
	}

	struct mie_select_chain_end *group_end = malloc(sizeof *group_end);
	if (!group_end) {
		return MIE_ERR_NO_MEMORY;
	}

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

	mie_select_node_get_value(group, chain_type, 0, &group_end->c_value);
	b_queue_push_back(&builder->b_graph->g_chain_ends, &group_end->c_entry);
	*out = group_end->c_value;

	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_entry *entry = b_queue_first(&builder->b_graph->g_nodes);
	while (entry) {
		struct mie_select_node *node
			= b_unbox(struct mie_select_node, entry, n_entry);

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

		entry = b_queue_next(entry);
	}

	return NULL;
}