#include "codegen.h"

#include <ivy/lang/codegen.h>
#include <mie/builder.h>
#include <mie/module.h>
#include <mie/value.h>
#include <stdlib.h>
#include <string.h>

enum ivy_status codegen_push_generator(
	struct ivy_codegen *gen, enum code_generator_type gen_type, void *arg)
{
	const struct code_generator *generator = get_code_generator(gen_type);
	if (!generator) {
		return IVY_ERR_INVALID_VALUE;
	}

	if (generator->g_state_size < sizeof(struct code_generator_state)) {
		return IVY_ERR_INTERNAL_FAILURE;
	}

	struct code_generator_state *state = malloc(generator->g_state_size);
	if (!state) {
		return IVY_ERR_NO_MEMORY;
	}

	memset(state, 0x0, generator->g_state_size);

	state->s_gen = generator;
	state->s_arg = arg;

	enum ivy_status status = IVY_OK;

	if (generator->g_state_init) {
		status = generator->g_state_init(gen, state);
	}

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

	b_queue_push_back(&gen->c_state, &state->s_entry);

	return IVY_OK;
}

enum ivy_status codegen_pop_generator(struct ivy_codegen *gen)
{
	b_queue_entry *entry = b_queue_pop_back(&gen->c_state);
	if (!entry) {
		return IVY_ERR_BAD_STATE;
	}

	struct code_generator_state *state
		= b_unbox(struct code_generator_state, entry, s_entry);

	enum ivy_status status = IVY_OK;

	if (state->s_gen->g_state_fini) {
		status = state->s_gen->g_state_fini(gen, state);
	}

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

	free(state);
	return IVY_OK;
}

void codegen_push_value(struct ivy_codegen *gen, struct mie_value *value)
{
	b_queue_push_back(&gen->c_ir_values, &value->v_entry);
}

struct mie_value *codegen_pop_value(struct ivy_codegen *gen)
{
	b_queue_entry *entry = b_queue_pop_back(&gen->c_ir_values);
	if (!entry) {
		return NULL;
	}

	return b_unbox(struct mie_value, entry, v_entry);
}

static struct code_generator_state *get_current_generator_state(
	struct ivy_codegen *gen)
{
	b_queue_entry *entry = b_queue_last(&gen->c_state);
	if (!entry) {
		return NULL;
	}

	return b_unbox(struct code_generator_state, entry, s_entry);
}

enum ivy_status ivy_codegen_create(struct ivy_codegen **out)
{
	struct ivy_codegen *gen = malloc(sizeof *gen);
	if (!gen) {
		return IVY_ERR_NO_MEMORY;
	}

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

	gen->c_ctx = mie_ctx_create();

	codegen_push_generator(gen, CODE_GENERATOR_UNIT, NULL);

	*out = gen;
	return IVY_OK;
}

void ivy_codegen_destroy(struct ivy_codegen *gen)
{
	if (gen->c_module) {
		mie_value_destroy(MIE_VALUE(gen->c_module));
	}

	if (gen->c_builder) {
		mie_builder_destroy(gen->c_builder);
	}

	if (gen->c_ctx) {
		mie_ctx_destroy(gen->c_ctx);
	}

	free(gen);
}

enum ivy_status ivy_codegen_start_module(struct ivy_codegen *gen)
{
	if (gen->c_module || gen->c_builder) {
		return IVY_ERR_BAD_STATE;
	}

	gen->c_module = mie_module_create();
	gen->c_builder = mie_builder_create(gen->c_ctx, gen->c_module);

	return IVY_OK;
}

enum ivy_status ivy_codegen_end_module(
	struct ivy_codegen *gen, struct mie_module **out)
{
	if (!gen->c_module) {
		return IVY_ERR_BAD_STATE;
	}

	struct mie_module *module = gen->c_module;
	mie_builder_destroy(gen->c_builder);

	gen->c_module = NULL;
	gen->c_builder = NULL;

	*out = module;

	return IVY_OK;
}

struct mie_module *ivy_codegen_get_current_module(struct ivy_codegen *gen)
{
	return gen->c_module;
}

enum ivy_status ivy_codegen_push_node(
	struct ivy_codegen *gen, struct ivy_ast_node *node,
	enum ivy_ast_iteration_type node_type)
{
	if (!gen->c_builder) {
		return IVY_ERR_BAD_STATE;
	}

	while (true) {
		struct code_generator_state *state
			= get_current_generator_state(gen);

		if (!state) {
			return IVY_ERR_BAD_STATE;
		}

		const struct code_generator *generator = state->s_gen;

		code_generator_node_callback func;
		switch (node_type) {
		case IVY_AST_ITERATION_PRE:
			func = generator->g_node_generators_pre[node->n_type];
			break;
		case IVY_AST_ITERATION_POST:
			func = generator->g_node_generators_post[node->n_type];
			break;
		default:
			return IVY_ERR_INVALID_VALUE;
		}

		if (!func) {
			return IVY_ERR_INVALID_VALUE;
		}

		struct code_generator_result result = func(gen, state, node);
		if (result.r_flags & CODEGEN_REPEAT_NODE) {
			continue;
		}

		return result.r_status;
	}

	return IVY_OK;
}