#include "../debug.h"
#include "codegen.h"

#include <assert.h>
#include <mie/ir/block.h>
#include <mie/ir/func.h>
#include <mie/ir/phi.h>
#include <mie/ir/value.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

struct cond_group_codegen_state {
	struct code_generator_state s_base;
	struct mie_block *s_end;
	b_queue s_edges;
};

static struct code_generator_result gen_cond_group(
	struct ivy_codegen *gen, struct code_generator_state *state,
	struct ivy_ast_node *node, size_t depth)
{
	struct cond_group_codegen_state *group
		= (struct cond_group_codegen_state *)state;

	return CODEGEN_RESULT_OK(0);
}

static struct code_generator_result gen_cond(
	struct ivy_codegen *gen, struct code_generator_state *state,
	struct ivy_ast_node *node, size_t depth)
{
	struct cond_group_codegen_state *group
		= (struct cond_group_codegen_state *)state;

	codegen_push_generator(gen, CODE_GENERATOR_COND, 0, group->s_end);
	return CODEGEN_RESULT_OK(CODEGEN_R_REPEAT_NODE);
}

static enum ivy_status state_init(
	struct ivy_codegen *gen, struct code_generator_state *state,
	uintptr_t argv, void *argp)
{
	debug_printf("codegen: start of cond group\n");
	struct cond_group_codegen_state *group
		= (struct cond_group_codegen_state *)state;
	struct mie_func *func = mie_builder_get_current_func(gen->c_builder);
	group->s_end = mie_func_create_block(func, "if.end");
	return IVY_OK;
}

static enum ivy_status get_eval_type(
	struct ivy_codegen *gen, struct cond_group_codegen_state *state,
	struct mie_type **out)
{
	struct mie_type *type = NULL;

	b_queue_iterator it;
	b_queue_foreach (&it, &state->s_edges) {
		struct mie_phi_edge *edge
			= b_unbox(struct mie_phi_edge, it.entry, e_entry);
		if (!edge->e_value) {
			continue;
		}

		struct mie_type *edge_type
			= mie_value_get_type(edge->e_value, gen->c_ctx);

		if (!type) {
			type = edge_type;
			continue;
		}
	}

	*out = type;
	return IVY_OK;
}

static enum ivy_status emit_phi_instr(
	struct ivy_codegen *gen, struct cond_group_codegen_state *state,
	struct mie_type *type, struct mie_value **result)
{
	size_t nr_edges = b_queue_length(&state->s_edges);
	if (!nr_edges) {
		return IVY_OK;
	}

	struct mie_phi_edge *edges = calloc(nr_edges, sizeof *edges);
	if (!edges) {
		return IVY_ERR_NO_MEMORY;
	}

	size_t i = 0;

	b_queue_iterator it;
	b_queue_foreach (&it, &state->s_edges) {
		struct mie_phi_edge *edge
			= b_unbox(struct mie_phi_edge, it.entry, e_entry);

		memcpy(&edges[i], edge, sizeof *edge);
		i++;
	}

	struct mie_value *phi = mie_builder_phi(
		gen->c_builder, type, edges, nr_edges, "if.result");
	*result = phi;

	return IVY_OK;
}

static enum ivy_status state_fini(
	struct ivy_codegen *gen, struct code_generator_state *state,
	struct code_generator_value *result)
{
	debug_printf("codegen: end of cond group\n");
	struct cond_group_codegen_state *group
		= (struct cond_group_codegen_state *)state;
	struct mie_func *func = mie_builder_get_current_func(gen->c_builder);
	struct mie_block *block = mie_builder_get_current_block(gen->c_builder);

	struct mie_type *result_type = NULL;
	enum ivy_status status = get_eval_type(gen, group, &result_type);
	if (status != IVY_OK) {
		return status;
	}

	if (!block->b_terminator) {
		mie_builder_br(gen->c_builder, group->s_end);
	}

	mie_func_insert_block(func, group->s_end, NULL);
	mie_builder_set_insert_point(gen->c_builder, group->s_end);

	struct mie_value *phi = NULL;
	if (result_type) {
		status = emit_phi_instr(gen, group, result_type, &phi);
	}

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

	if (phi) {
		code_generator_value_set_mie_value(result, phi);
	}

	return IVY_OK;
}

static struct code_generator_result value_received(
	struct ivy_codegen *gen, struct code_generator_state *state,
	struct code_generator_value *value)
{
	struct cond_group_codegen_state *cond
		= (struct cond_group_codegen_state *)state;
	assert(value->v_type == CODE_GENERATOR_VALUE_PHI_EDGE);

	struct mie_phi_edge *edge = code_generator_value_get_phi_edge(value);
	b_queue_push_back(&cond->s_edges, &edge->e_entry);

	return CODEGEN_RESULT_OK(0);
}

struct code_generator cond_group_generator = {
	.g_type = CODE_GENERATOR_COND_GROUP,
	.g_state_size = sizeof(struct cond_group_codegen_state),
	.g_state_init = state_init,
	.g_state_fini = state_fini,
	.g_value_received = value_received,
	.g_node_generators = {
		NODE_CODEGEN(COND_GROUP, gen_cond_group),
		NODE_CODEGEN(COND, gen_cond),
	},
};