#include "mie/select/opcode.h"

#include <ivy/asm/mie.h>
#include <ivy/opcode.h>
#include <mie/ctx.h>
#include <mie/ir/const.h>
#include <mie/ir/msg.h>
#include <mie/select/builder.h>
#include <mie/select/graph.h>
#include <mie/select/node.h>
#include <mie/target/select.h>
#include <stdio.h>
#include <stdlib.h>

#define NODE_NAME(op)                                                          \
	case IVY_SELECT_OP_##op:                                               \
		return snprintf(out, max, #op)

#define NODE_NAME_(op, name)                                                   \
	case IVY_SELECT_OP_##op:                                               \
		return snprintf(out, max, name)

static size_t node_name(
	const struct mie_target *target, unsigned int opcode, char *out, size_t max)
{
	switch (opcode) {
		NODE_NAME(LDR);
		NODE_NAME(STR);
		NODE_NAME(PUSH);
		NODE_NAME(POP);
		NODE_NAME(MSG);
		NODE_NAME(ADD);
		NODE_NAME(SUB);
		NODE_NAME(MUL);
		NODE_NAME(DIV);
		NODE_NAME_(C_EQ, "C.EQ");
		NODE_NAME_(C_NE, "C.NE");
		NODE_NAME_(C_LT, "C.LT");
		NODE_NAME_(C_LE, "C.LE");
		NODE_NAME_(C_GT, "C.GT");
		NODE_NAME_(C_GE, "C.GE");
		NODE_NAME(BR);
		NODE_NAME_(BR_T, "BR.T");
		NODE_NAME_(BR_F, "BR.F");
		NODE_NAME_(OB_C, "OB.C");
		NODE_NAME_(OB_E, "OB.E");
		NODE_NAME_(LAM_C, "LAM.C");
		NODE_NAME_(IT_G, "IT.G");
		NODE_NAME_(IT_N, "IT.N");
		NODE_NAME_(IT_V, "IT.V");
		NODE_NAME_(STK_A, "STK.A");
		NODE_NAME_(STK_F, "STK.F");
		NODE_NAME(RET);
		NODE_NAME_(RET_N, "RET.N");
		NODE_NAME_(REGISTER, "Register");
		NODE_NAME_(SELECTOR, "Selector");
	default:
		return snprintf(out, max, "UNKNOWN");
	}

	return 0;
}

static enum mie_status get_sp_register(
	const struct mie_target *target, struct mie_select_builder *builder,
	struct mie_select_value *out)
{
	struct mie_ctx *ctx = mie_select_builder_get_ctx(builder);
	struct mie_type *ptr_type = mie_ctx_get_type(ctx, MIE_TYPE_PTR);

	struct mie_select_node *sp_node = mie_select_builder_find_node_with_ivalue(
		builder, target, IVY_SELECT_OP_REGISTER, IVY_SELECT_REG_SP);

	if (sp_node) {
		mie_select_node_get_value(sp_node, ptr_type, 0, out);
		return MIE_SUCCESS;
	}

	struct mie_select_graph *graph = mie_select_builder_get_graph(builder);
	enum mie_status status = mie_select_graph_get_node(
		graph, target, IVY_SELECT_OP_REGISTER, NULL, 0, &ptr_type, 1,
		&sp_node);
	if (status != MIE_SUCCESS) {
		return status;
	}

	sp_node->n_value.i = IVY_SELECT_REG_SP;
	mie_select_node_set_description(sp_node, "$sp");
	mie_select_node_get_value(sp_node, ptr_type, 0, out);
	return MIE_SUCCESS;
}

static enum mie_status get_selector(
	const struct mie_target *target, struct mie_select_builder *builder,
	struct mie_selector *sel, struct mie_select_value *out)
{
	struct mie_ctx *ctx = mie_select_builder_get_ctx(builder);
	struct mie_type *ptr_type = mie_ctx_get_type(ctx, MIE_TYPE_PTR);
	struct mie_select_graph *graph = mie_select_builder_get_graph(builder);

	b_queue *nodes = &graph->g_nodes;
	b_queue_entry *entry = b_queue_first(nodes);
	while (entry) {
		struct mie_select_node *node
			= b_unbox(struct mie_select_node, entry, n_entry);

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

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

		if (!node->n_value.v || !mie_value_is_selector(node->n_value.v)) {
			goto skip;
		}

		struct mie_selector *sel_node = (struct mie_selector *)node;
		if (strcmp(sel->sel_value, sel_node->sel_value) != 0) {
			goto skip;
		}

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

	struct mie_select_node *sel_node;
	enum mie_status status = mie_select_graph_get_node(
		graph, target, IVY_SELECT_OP_SELECTOR, NULL, 0, &ptr_type, 1,
		&sel_node);
	if (status != MIE_SUCCESS) {
		return status;
	}

	sel_node->n_value.v = MIE_VALUE(sel);
	mie_select_node_set_description(sel_node, "%s", sel->sel_value);
	mie_select_node_get_value(sel_node, ptr_type, 0, out);

	return MIE_SUCCESS;
}

static enum mie_status load_selector(
	const struct mie_target *target, struct mie_select_builder *builder,
	struct mie_select_value *sel, struct mie_select_value *out)
{
	struct mie_ctx *ctx = mie_select_builder_get_ctx(builder);
	struct mie_type *id_type = mie_ctx_get_type(ctx, MIE_TYPE_ID);
	struct mie_select_graph *graph = mie_select_builder_get_graph(builder);
	struct mie_select_node *load_sel;

	struct mie_select_value *operands[] = {
		&graph->g_entry,
		sel,
	};
	const size_t nr_operands = sizeof operands / sizeof operands[0];

	struct mie_type *results[] = {
		id_type,
	};
	const size_t nr_results = sizeof results / sizeof results[0];

	enum mie_status status = mie_select_graph_get_node(
		graph, mie_target_builtin(), MIE_SELECT_OP_LOAD, operands,
		nr_operands, results, nr_results, &load_sel);
	if (status != MIE_SUCCESS) {
		return status;
	}

	mie_select_node_get_value(load_sel, id_type, 0, out);
	return MIE_SUCCESS;
}

static enum mie_status create_stack_alloc(
	const struct mie_target *target, struct mie_select_builder *builder,
	size_t count, struct mie_select_value *incoming_chain,
	struct mie_select_value *out)
{
	struct mie_select_node *alloc;
	enum mie_status status;

	struct mie_ctx *ctx = mie_select_builder_get_ctx(builder);
	struct mie_type *chain = mie_ctx_get_type(ctx, MIE_TYPE_OTHER);
	struct mie_type *i32 = mie_ctx_get_int_type(ctx, 32);

	struct mie_select_value count_value;
	mie_select_builder_get_const(builder, count, &count_value);

	struct mie_select_value *operands[] = {
		&count_value,
		incoming_chain,
	};
	const size_t nr_operands = sizeof operands / sizeof operands[0];

	struct mie_select_graph *graph = mie_select_builder_get_graph(builder);
	status = mie_select_graph_get_node(
		graph, target, IVY_SELECT_OP_STK_A, operands, nr_operands,
		&chain, 1, &alloc);
	if (status != MIE_SUCCESS) {
		return status;
	}

	mie_select_node_get_value(alloc, chain, 0, out);
	mie_select_node_set_description(
		alloc, "(allocate %zu slots on stack)", count);

	return MIE_SUCCESS;
}

static enum mie_status create_stack_free(
	const struct mie_target *target, struct mie_select_builder *builder,
	size_t count, struct mie_select_value *incoming_chain,
	struct mie_select_value *out)
{
	struct mie_select_node *free_node;
	enum mie_status status;

	struct mie_ctx *ctx = mie_select_builder_get_ctx(builder);
	struct mie_type *chain = mie_ctx_get_type(ctx, MIE_TYPE_OTHER);
	struct mie_type *i32 = mie_ctx_get_int_type(ctx, 32);

	struct mie_select_value count_value;
	mie_select_builder_get_const(builder, count, &count_value);

	struct mie_select_value *operands[] = {
		&count_value,
		incoming_chain,
	};
	const size_t nr_operands = sizeof operands / sizeof operands[0];

	struct mie_select_graph *graph = mie_select_builder_get_graph(builder);
	status = mie_select_graph_get_node(
		graph, target, IVY_SELECT_OP_STK_F, operands, nr_operands,
		&chain, 1, &free_node);
	if (status != MIE_SUCCESS) {
		return status;
	}

	mie_select_node_set_description(
		free_node, "(free %zu slots on stack)", count);
	mie_select_node_get_value(free_node, chain, 0, out);

	return MIE_SUCCESS;
}

static enum mie_status create_chain_group(
	const struct mie_target *target, struct mie_select_builder *builder,
	struct mie_select_value **incoming_chains, size_t count,
	struct mie_select_value *out)
{
	struct mie_select_graph *graph = mie_select_builder_get_graph(builder);
	struct mie_ctx *ctx = mie_select_builder_get_ctx(builder);
	struct mie_select_node *group = NULL;

	struct mie_type *chain_type = mie_ctx_get_type(ctx, MIE_TYPE_OTHER);

	enum mie_status status = mie_select_graph_get_node(
		graph, mie_target_builtin(), MIE_SELECT_OP_CHAIN_GROUP,
		incoming_chains, count, &chain_type, 1, &group);
	if (status != MIE_SUCCESS) {
		return status;
	}

	mie_select_node_get_value(group, chain_type, 0, out);
	return MIE_SUCCESS;
}

static enum mie_status put_param_on_stack(
	const struct mie_target *target, struct mie_select_builder *builder,
	struct mie_select_value *in_chain, struct mie_select_value *param,
	size_t offset, struct mie_select_value *out_chain)
{
	struct mie_select_value sp;
	struct mie_select_value offset_value;
	struct mie_type *chain_type;

	struct mie_select_graph *graph = mie_select_builder_get_graph(builder);
	struct mie_ctx *ctx = mie_select_builder_get_ctx(builder);

	get_sp_register(target, builder, &sp);
	mie_select_builder_get_const(builder, offset, &offset_value);
	chain_type = mie_ctx_get_type(ctx, MIE_TYPE_OTHER);

	struct mie_select_value *operands[] = {
		in_chain,
		param,
		&sp,
		&offset_value,
	};
	const size_t nr_operands = sizeof operands / sizeof operands[0];

	struct mie_select_node *store;
	enum mie_status status = mie_select_graph_get_node(
		graph, target, IVY_SELECT_OP_STR, operands, nr_operands,
		&chain_type, 1, &store);

	mie_select_node_set_description(
		store, "(store N%zu on stack + %zu)", param->v_node->n_id, offset);
	mie_select_node_get_value(store, chain_type, 0, out_chain);
	return MIE_SUCCESS;
}

static enum mie_status lower_msg(
	const struct mie_target *target, struct mie_select_builder *builder,
	struct mie_msg *msg, struct mie_select_value *result)
{
	enum mie_status status = MIE_SUCCESS;
	struct mie_select_graph *graph = mie_select_builder_get_graph(builder);
	struct mie_ctx *ctx = mie_select_builder_get_ctx(builder);

	struct mie_select_value **param_chains = calloc(
		b_max(size_t, 2, msg->msg_nr_args + 1),
		sizeof(struct mie_select_value *));
	size_t nr_param_chains = 0;
	struct mie_select_value *chain, *recipient, nr_args;

	chain = mie_select_builder_get_mem_access(builder, msg->msg_recipient);
	recipient = mie_select_builder_get_value(builder, msg->msg_recipient);
	mie_select_builder_get_const(builder, msg->msg_nr_args, &nr_args);

	if (chain) {
		param_chains[nr_param_chains++] = chain;
	}

	for (size_t i = 0; i < msg->msg_nr_args; i++) {
		struct mie_select_value *chain = mie_select_builder_get_mem_access(
			builder, msg->msg_args[i]);
		if (chain) {
			param_chains[nr_param_chains++] = chain;
		}
	}

	if (nr_param_chains == 0) {
		param_chains[nr_param_chains++] = &graph->g_entry;
	}

	struct mie_select_value pre_arg_init_chain, post_arg_init_chain;
	if (nr_param_chains > 1) {
		status = create_chain_group(
			target, builder, param_chains, nr_param_chains,
			&pre_arg_init_chain);
	} else {
		pre_arg_init_chain = *param_chains[0];
	}

	struct mie_select_value stack_alloc, stack_free;
	status = create_stack_alloc(
		target, builder, msg->msg_nr_args, &pre_arg_init_chain,
		&stack_alloc);
	if (status != MIE_SUCCESS) {
		return status;
	}

	struct mie_select_value *stack_chains = calloc(
		b_max(size_t, 1, msg->msg_nr_args),
		sizeof(struct mie_select_value));
	size_t nr_stack_chains = 0;
	nr_param_chains = 0;

	for (size_t i = 0; i < msg->msg_nr_args; i++) {
		struct mie_select_value *param_value
			= mie_select_builder_get_value(builder, msg->msg_args[i]);
		status = put_param_on_stack(
			target, builder, &stack_alloc, param_value, i,
			&stack_chains[nr_stack_chains]);
		param_chains[nr_param_chains] = &stack_chains[nr_stack_chains];

		nr_param_chains++;
		nr_stack_chains++;
	}

	if (nr_param_chains > 1) {
		status = create_chain_group(
			target, builder, param_chains, nr_param_chains,
			&post_arg_init_chain);
	} else {
		post_arg_init_chain = *param_chains[0];
	}

	free(param_chains);
	free(stack_chains);

	struct mie_select_value sel, load_sel;
	get_selector(target, builder, MIE_SELECTOR(msg->msg_selector), &sel);
	load_selector(target, builder, &sel, &load_sel);

	struct mie_select_value *msg_params[] = {
		&post_arg_init_chain,
		recipient,
		&load_sel,
		&nr_args,
	};
	const size_t nr_msg_params = sizeof msg_params / sizeof msg_params[0];

	struct mie_type *msg_results[] = {
		mie_ctx_get_type(ctx, MIE_TYPE_OTHER),
		mie_ctx_get_type(ctx, MIE_TYPE_GLUE),
	};
	const size_t nr_msg_results = sizeof msg_results / sizeof msg_results[0];

	struct mie_select_node *msg_node;
	status = mie_select_graph_get_node(
		graph, target, IVY_SELECT_OP_MSG, msg_params, nr_msg_params,
		msg_results, nr_msg_results, &msg_node);

	result->v_node = msg_node;
	result->v_index = 0;

	return MIE_SUCCESS;
}

const struct mie_target_select_ops __ivy_select_ops = {
	.s_node_name = node_name,
	.s_lower_msg = lower_msg,
};