#include "codegen.h"

#include <mie/ir/block.h>
#include <mie/ir/func.h>
#include <mie/ir/module.h>

struct while_codegen_state {
	struct code_generator_state s_base;
	struct mie_value *s_cond_value;

	struct {
		struct mie_block *b_cond;
		struct mie_block *b_body;
		struct mie_block *b_end;
	} s_blocks;
};

static struct code_generator_result gen_expr(
	struct ivy_codegen *gen, struct code_generator_state *state,
	struct ivy_ast_node *node, size_t depth)
{
	struct while_codegen_state *while_loop
		= (struct while_codegen_state *)state;

	if (!while_loop->s_cond_value) {
		mie_builder_br(gen->c_builder, while_loop->s_blocks.b_cond);

		struct mie_func *func
			= mie_builder_get_current_func(gen->c_builder);
		mie_func_insert_block(func, while_loop->s_blocks.b_cond, NULL);
		mie_builder_set_insert_point(
			gen->c_builder, while_loop->s_blocks.b_cond);
	}

	codegen_push_generator(gen, CODE_GENERATOR_EXPR, 0, NULL);
	return CODEGEN_RESULT_OK(CODEGEN_R_REPEAT_NODE);
}

static struct code_generator_result gen_while_loop(
	struct ivy_codegen *gen, struct code_generator_state *state,
	struct ivy_ast_node *node, size_t depth)
{
	return CODEGEN_RESULT_OK(0);
}

static struct code_generator_result gen_block(
	struct ivy_codegen *gen, struct code_generator_state *state,
	struct ivy_ast_node *node, size_t depth)
{
	struct while_codegen_state *while_loop
		= (struct while_codegen_state *)state;
	struct mie_func *func = mie_builder_get_current_func(gen->c_builder);

	// mie_builder_br(gen->c_builder, while_loop->s_blocks.b_cond);
	//  mie_func_insert_block(func, while_loop->s_blocks.b_cond, NULL);
	//  mie_builder_set_insert_point(gen->c_builder, while_loop->s_blocks.b_cond);

	codegen_push_generator(gen, CODE_GENERATOR_BLOCK, 0, NULL);
	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)
{
	struct while_codegen_state *while_loop
		= (struct while_codegen_state *)state;
	struct mie_func *func = mie_builder_get_current_func(gen->c_builder);

	while_loop->s_blocks.b_cond = mie_func_create_block(func, "while.cond");
	while_loop->s_blocks.b_body = mie_func_create_block(func, "while.body");
	while_loop->s_blocks.b_end = mie_func_create_block(func, "while.end");

	state->s_loop_break_target = while_loop->s_blocks.b_end;
	state->s_loop_repeat_target = while_loop->s_blocks.b_cond;

	return IVY_OK;
}

static enum ivy_status gen_cond_branch(
	struct ivy_codegen *gen, struct while_codegen_state *while_loop)
{
	struct mie_func *func = mie_builder_get_current_func(gen->c_builder);

	mie_builder_br_if(
		gen->c_builder, while_loop->s_cond_value,
		while_loop->s_blocks.b_body, while_loop->s_blocks.b_end);

	return IVY_OK;
}

static enum ivy_status gen_body(
	struct ivy_codegen *gen, struct while_codegen_state *while_loop)
{
	struct mie_func *func = mie_builder_get_current_func(gen->c_builder);
	mie_func_insert_block(func, while_loop->s_blocks.b_body, NULL);
	mie_builder_set_insert_point(gen->c_builder, while_loop->s_blocks.b_body);

	return IVY_OK;
}

static enum ivy_status gen_end_block(
	struct ivy_codegen *gen, struct while_codegen_state *while_loop)
{
	mie_builder_br(gen->c_builder, while_loop->s_blocks.b_cond);

	struct mie_func *func = mie_builder_get_current_func(gen->c_builder);
	mie_func_insert_block(func, while_loop->s_blocks.b_end, NULL);
	mie_builder_set_insert_point(gen->c_builder, while_loop->s_blocks.b_end);

	return IVY_OK;
}

static enum ivy_status state_fini(
	struct ivy_codegen *gen, struct code_generator_state *state,
	struct code_generator_value *result)
{
	struct while_codegen_state *while_loop
		= (struct while_codegen_state *)state;
	gen_end_block(gen, while_loop);
	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 while_codegen_state *while_loop
		= (struct while_codegen_state *)state;
	struct mie_func *func = mie_builder_get_current_func(gen->c_builder);

	if (!while_loop->s_cond_value) {
		while_loop->s_cond_value
			= code_generator_value_get_mie_value(value);
		gen_cond_branch(gen, while_loop);
		gen_body(gen, while_loop);
		return CODEGEN_RESULT_OK(0);
	}

	return CODEGEN_RESULT_OK(0);
}

struct code_generator while_loop_generator = {
	.g_type = CODE_GENERATOR_WHILE_LOOP,
	.g_state_size = sizeof(struct while_codegen_state),
	.g_state_init = state_init,
	.g_state_fini = state_fini,
	.g_value_received = value_received,
	.g_expr_generator = gen_expr,
	.g_node_generators = {
		NODE_CODEGEN(WHILE_LOOP, gen_while_loop),
		NODE_CODEGEN(BLOCK, gen_block),
	},
};