#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>

/* new codegen layout:
 *   - traverse the AST in pre-order (i.e. in the order it's displayed by debug)
 *   - each node type (var, msg, expr) gets its own code generator type
 *   - expr code generator is limited to simple operator arithmetic.
 *   - when a code generator encounters a node of an equal or lower depth than
 *     the node that started it, it has reached the end of its subtree.
 *     - depth is supplied by the ast iterator.
 *   - need to devise a way for expression code generator to "return" a
 * mie_value for its parent generator to use.
 */

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, struct mie_value **result)
{
	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, result);
	}

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

	free(state);
	return IVY_OK;
}

static struct codegen_var *code_generator_scope_get_variable(
	struct code_generator_scope *scope, const char *ident)
{
	b_queue_iterator it;
	b_queue_foreach (&it, &scope->s_vars) {
		struct codegen_var *var
			= b_unbox(struct codegen_var, it.entry, v_entry);

		if (!strcmp(var->v_ident->n_content->t_str, ident)) {
			return var;
		}
	}
	return NULL;
}

static bool code_generator_scope_is_top_level(struct code_generator_scope *scope)
{
	switch (scope->s_type) {
	case CODE_GENERATOR_SCOPE_FUNC:
	case CODE_GENERATOR_SCOPE_UNIT:
		return true;
	default:
		return false;
	}
}

static void code_generator_scope_destroy(struct code_generator_scope *scope)
{
	/* TODO clean up variables */
	free(scope);
}

struct code_generator_scope *codegen_push_scope(struct ivy_codegen *gen)
{
	struct code_generator_scope *scope = malloc(sizeof *scope);
	if (!scope) {
		return NULL;
	}

	memset(scope, 0x0, sizeof *scope);
	b_queue_push_back(&gen->c_scope, &scope->s_entry);

	return scope;
}

void codegen_pop_scope(struct ivy_codegen *gen)
{
	b_queue_entry *entry = b_queue_pop_back(&gen->c_scope);
	if (!entry) {
		return;
	}

	struct code_generator_scope *scope
		= b_unbox(struct code_generator_scope, entry, s_entry);
	code_generator_scope_destroy(scope);
}

struct code_generator_scope *codegen_get_current_scope(struct ivy_codegen *gen)
{
	b_queue_entry *entry = b_queue_last(&gen->c_scope);
	if (!entry) {
		return NULL;
	}

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

struct codegen_value *codegen_value_create(
	struct ivy_ast_node *node, struct mie_value *val)
{
	struct codegen_value *out = malloc(sizeof *out);
	if (!out) {
		return NULL;
	}

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

	out->v_node = node;
	out->v_value = val;

	return out;
}

void codegen_value_destroy(struct codegen_value *val)
{
	free(val);
}

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

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

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

struct codegen_var *codegen_resolve_variable(
	struct ivy_codegen *gen, const char *ident)
{
	b_queue_entry *cur = b_queue_last(&gen->c_scope);
	while (cur) {
		struct code_generator_scope *scope
			= b_unbox(struct code_generator_scope, cur, s_entry);
		struct codegen_var *var
			= code_generator_scope_get_variable(scope, ident);

		if (var) {
			return var;
		}

		if (code_generator_scope_is_top_level(scope)) {
			break;
		}
	}

	return NULL;
}

enum ivy_status code_generator_scope_put_variable(
	struct code_generator_scope *scope, struct ivy_ast_ident_node *ident,
	struct mie_value *ptr)
{
	struct codegen_var *var = malloc(sizeof *var);
	if (!var) {
		return IVY_ERR_NO_MEMORY;
	}

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

	var->v_ident = ident;
	var->v_ptr = ptr;

	b_queue_push_back(&scope->s_vars, &var->v_entry);

	return IVY_OK;
}

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();

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

static enum ivy_status pop_generator_recurse(struct ivy_codegen *gen)
{
	while (1) {
		struct mie_value *value = NULL;
		codegen_pop_generator(gen, &value);
		struct code_generator_state *state
			= get_current_generator_state(gen);

		if (!state || !state->s_gen->g_value_received) {
			return IVY_OK;
		}

		struct code_generator_result result
			= state->s_gen->g_value_received(gen, state, value);

		if (!(result.r_flags & CODEGEN_POP_GENERATOR)) {
			return result.r_status;
		}
	}

	return IVY_OK;
}

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

	if (b_queue_empty(&gen->c_state)) {
		const struct code_generator *generator
			= get_root_code_generator(node->n_type);
		if (!generator) {
			return IVY_ERR_BAD_SYNTAX;
		}

		codegen_push_generator(gen, generator->g_type, 0);
	}

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

		if (!state) {
			return IVY_ERR_BAD_STATE;
		}

		enum ivy_status status = IVY_OK;
		struct code_generator_result result = {};

		if (!state->s_root) {
			state->s_root = node;
			state->s_depth = node_depth;
		} else if (node_depth <= state->s_depth) {
			status = pop_generator_recurse(gen);
		}

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

		state = get_current_generator_state(gen);
		const struct code_generator *generator = state->s_gen;

		code_generator_node_callback func
			= generator->g_node_generators[node->n_type];

		if (!func) {
			return IVY_OK;
		}

		result = func(gen, state, node, node_depth);

		if (result.r_flags & CODEGEN_REPEAT_NODE) {
			continue;
		}

		return result.r_status;
	}

	return IVY_OK;
}

enum ivy_status ivy_codegen_push_eof(struct ivy_codegen *gen)
{
	struct code_generator_result result;

	while (!b_queue_empty(&gen->c_state)) {
		struct mie_value *value = NULL;
		codegen_pop_generator(gen, &value);
		struct code_generator_state *state
			= get_current_generator_state(gen);

		if (!state || !state->s_gen->g_value_received) {
			continue;
		}

		result = state->s_gen->g_value_received(gen, state, value);

		if (result.r_status != IVY_OK) {
			break;
		}
	}

	return result.r_status;
}