#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,
	uintptr_t argv, void *argp)
{
	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;

	enum ivy_status status = IVY_OK;

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

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

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

enum ivy_status codegen_define_variable(
	struct ivy_codegen *gen, const char *ident, const struct codegen_var *var)
{
	b_queue_entry *entry = b_queue_last(&gen->c_state);
	if (!entry) {
		return IVY_ERR_NO_ENTRY;
	}

	bool resolved = false;

	while (entry) {
		struct code_generator_state *state
			= b_unbox(struct code_generator_state, entry, s_entry);
		const struct code_generator *generator = state->s_gen;
		enum ivy_status status = IVY_ERR_NO_ENTRY;

		if (generator->g_define_var) {
			status = generator->g_define_var(gen, state, ident, var);
		}

		if (status == IVY_OK) {
			break;
		}

		entry = b_queue_prev(entry);
	}

	return IVY_ERR_NOT_SUPPORTED;
}

enum ivy_status resolve_global_variable(
	struct ivy_codegen *gen, const char *ident, struct codegen_var *out)
{
	struct mie_value *var_ptr
		= mie_builder_get_data_ptr(gen->c_builder, ident);

	out->v_type = mie_ctx_get_type(gen->c_ctx, MIE_TYPE_ID);
	out->v_node = NULL;
	out->v_flags = CODEGEN_VAR_F_PTR;
	out->v_value = var_ptr;

	return IVY_OK;
}

enum ivy_status codegen_resolve_variable(
	struct ivy_codegen *gen, const char *ident, struct codegen_var *out)
{
	b_queue_entry *entry = b_queue_last(&gen->c_state);
	if (!entry) {
		return IVY_ERR_NO_ENTRY;
	}

	bool resolved = false;

	while (entry) {
		struct code_generator_state *state
			= b_unbox(struct code_generator_state, entry, s_entry);
		const struct code_generator *generator = state->s_gen;
		enum ivy_status status = IVY_ERR_NO_ENTRY;

		if (generator->g_resolve_var) {
			status = generator->g_resolve_var(gen, state, ident, out);
		}

		if (status == IVY_OK) {
			resolved = true;
			break;
		}

		entry = b_queue_prev(entry);
	}

	if (!resolved) {
		return resolve_global_variable(gen, ident, out);
	}

	if (entry) {
		entry = b_queue_next(entry);
	}

	while (entry) {
		struct code_generator_state *state
			= b_unbox(struct code_generator_state, entry, s_entry);
		const struct code_generator *generator = state->s_gen;
		enum ivy_status status = IVY_OK;

		if (generator->g_capture_var) {
			status = generator->g_capture_var(gen, state, ident, out);
		}

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

		entry = b_queue_next(entry);
	}

	return IVY_OK;
}

struct mie_value *codegen_load_variable(
	struct ivy_codegen *gen, struct codegen_var *var)
{
	return mie_builder_load(gen->c_builder, var->v_type, var->v_value, NULL);
}

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, size_t node_depth)
{
	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);

		bool should_pop = ((result.r_flags & CODEGEN_R_POP_GENERATOR) != 0)
		               || node_depth <= state->s_depth;

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

		if (!should_pop) {
			return IVY_OK;
		}
	}

	return IVY_OK;
}

static bool node_is_expr(struct ivy_ast_node *node)
{
	switch (node->n_type) {
	case IVY_AST_OP:
	case IVY_AST_MSG:
	case IVY_AST_LAMBDA:
	case IVY_AST_DISCARD:
	case IVY_AST_INT:
	case IVY_AST_DOUBLE:
	case IVY_AST_STRING:
	case IVY_AST_FSTRING:
	case IVY_AST_ATOM:
	case IVY_AST_IDENT:
	case IVY_AST_CASCADE:
	case IVY_AST_MATCH:
	case IVY_AST_COND_GROUP:
	case IVY_AST_TUPLE:
	case IVY_AST_PKG_STATIC:
	case IVY_AST_PKG_DYNAMIC:
		return true;
	default:
		return false;
	}
}

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, NULL);
	}

	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, node_depth);
		}

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

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

		bool is_expr = node_is_expr(node);
		code_generator_node_callback func
			= generator->g_node_generators[node->n_type];

		if (!func && is_expr) {
			func = generator->g_expr_generator;
		}

		if (!func) {
			return IVY_OK;
		}

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

		if (result.r_flags & CODEGEN_R_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;
}