#include "codegen.h"
#include "var-map.h"

#include <blue/core/stringstream.h>
#include <blue/object/list.h>
#include <mie/ir/block.h>
#include <mie/ir/func.h>
#include <mie/ir/module.h>
#include <mie/ir/record.h>
#include <stdlib.h>

enum {
	ENSURE_OPEN_BLOCK = 0x01u,
};

struct unit_codegen_state {
	struct code_generator_state s_base;
	struct code_generator_scope *s_scope;
	/* function for top-level statements */
	// struct mie_block *s_immediate_alloca;
	struct mie_func *s_immediate;
	struct codegen_var_map s_vars;
};

static enum ivy_status switch_to_immediate_func(
	struct ivy_codegen *gen, struct unit_codegen_state *unit, long flags)
{
	if (!unit->s_immediate) {
		struct mie_type *ret_type
			= mie_ctx_get_type(gen->c_ctx, MIE_TYPE_VOID);
		unit->s_immediate = mie_func_create(MIE_FUNC_STATIC, ret_type);

		if (!unit->s_immediate) {
			return IVY_ERR_NO_MEMORY;
		}

		mie_module_add_function(
			gen->c_module, unit->s_immediate, "init");
	}

	struct mie_block *block = mie_func_get_last_block(unit->s_immediate);
	if (!block) {
#if 0
		unit->s_immediate_alloca
			= mie_func_create_block(unit->s_immediate, "alloca");
		mie_func_insert_block(
			unit->s_immediate, unit->s_immediate_alloca, NULL);
#endif
		block = mie_func_create_block(unit->s_immediate, "entry");
		mie_func_insert_block(unit->s_immediate, block, NULL);
	}

	if (mie_block_is_terminated(block) && (flags & ENSURE_OPEN_BLOCK)) {
		block = mie_func_create_block(unit->s_immediate, NULL);
		mie_func_insert_block(unit->s_immediate, block, NULL);
	}

	mie_builder_set_insert_point(gen->c_builder, block);

	return IVY_OK;
}

#if 0
static enum ivy_status switch_to_immediate_alloca(
	struct ivy_codegen *gen, struct unit_codegen_state *unit)
{
	enum ivy_status status = switch_to_immediate_func(gen, unit);

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

	mie_builder_set_insert_point(gen->c_builder, unit->s_immediate_alloca);

	return IVY_OK;
}
#endif

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 unit_codegen_state *unit = (struct unit_codegen_state *)state;
	struct mie_builder *builder = gen->c_builder;
	struct mie_func *current = mie_builder_get_current_func(builder);

	enum ivy_status status = IVY_OK;
	if (!current || current != unit->s_immediate) {
		status = switch_to_immediate_func(gen, unit, ENSURE_OPEN_BLOCK);
	}

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

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

static struct code_generator_result gen_for_loop(
	struct ivy_codegen *gen, struct code_generator_state *state,
	struct ivy_ast_node *node, size_t depth)
{
	struct unit_codegen_state *unit = (struct unit_codegen_state *)state;
	struct mie_builder *builder = gen->c_builder;
	struct mie_func *current = mie_builder_get_current_func(builder);

	enum ivy_status status = IVY_OK;
	if (!current || current != unit->s_immediate) {
		status = switch_to_immediate_func(gen, unit, ENSURE_OPEN_BLOCK);
	}

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

	codegen_push_generator(
		gen, CODE_GENERATOR_FOR_LOOP, CODEGEN_F_IGNORE_RESULT, 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)
{
	struct unit_codegen_state *unit = (struct unit_codegen_state *)state;
	struct mie_builder *builder = gen->c_builder;
	struct mie_func *current = mie_builder_get_current_func(builder);

	enum ivy_status status = IVY_OK;
	if (!current || current != unit->s_immediate) {
		status = switch_to_immediate_func(gen, unit, ENSURE_OPEN_BLOCK);
	}

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

	codegen_push_generator(
		gen, CODE_GENERATOR_WHILE_LOOP, CODEGEN_F_IGNORE_RESULT, NULL);
	return CODEGEN_RESULT_OK(CODEGEN_R_REPEAT_NODE);
}

static void serialise_package_name(b_queue *parts, b_stringstream *out)
{
	b_stringstream_add(out, "_ZN");

	b_queue_iterator it;
	b_queue_foreach (&it, parts) {
		struct ivy_token *tok
			= b_unbox(struct ivy_token, it.entry, t_entry);
		size_t len = strlen(tok->t_str);
		b_stringstream_addf(out, "%zu%s", len, tok->t_str);
	}

	b_stringstream_add(out, "E");
}

static struct code_generator_result gen_unit_package(
	struct ivy_codegen *gen, struct code_generator_state *state,
	struct ivy_ast_node *node, size_t depth)
{
	struct unit_codegen_state *unit = (struct unit_codegen_state *)state;
	struct ivy_ast_unit_package_node *pkg
		= (struct ivy_ast_unit_package_node *)node;

	b_stringstream str;
	b_stringstream_begin_dynamic(&str);
	serialise_package_name(&pkg->n_ident, &str);
	char *ident = b_stringstream_end(&str);
	struct mie_value *ident_val = mie_ctx_get_string(gen->c_ctx, ident);
	free(ident);

	mie_builder_put_record(
		gen->c_builder, MIE_CONST(ident_val), "package_scope");

	return CODEGEN_RESULT_OK(0);
}

static struct code_generator_result gen_unit_import(
	struct ivy_codegen *gen, struct code_generator_state *state,
	struct ivy_ast_node *node, size_t depth)
{
	struct unit_codegen_state *unit = (struct unit_codegen_state *)state;
	struct ivy_ast_unit_package_node *pkg
		= (struct ivy_ast_unit_package_node *)node;

	struct mie_record *rec
		= mie_builder_get_record(gen->c_builder, "import");
	struct mie_array *import_list = NULL;

	if (!rec) {
		import_list = MIE_ARRAY(mie_ctx_create_array(gen->c_ctx));
		rec = mie_builder_put_record(
			gen->c_builder, MIE_CONST(import_list), "import");
	} else {
		import_list = MIE_ARRAY(rec->r_value);
	}

	b_stringstream str;
	b_stringstream_begin_dynamic(&str);
	serialise_package_name(&pkg->n_ident, &str);
	char *ident = b_stringstream_end(&str);
	struct mie_value *ident_val = mie_ctx_get_string(gen->c_ctx, ident);
	free(ident);

	b_list_push_back(import_list->a_values, ident_val);

	return CODEGEN_RESULT_OK(0);
}

static struct code_generator_result gen_return(
	struct ivy_codegen *gen, struct code_generator_state *state,
	struct ivy_ast_node *node, size_t depth)
{
	struct unit_codegen_state *unit = (struct unit_codegen_state *)state;
	struct mie_builder *builder = gen->c_builder;
	struct mie_func *current = mie_builder_get_current_func(builder);

	enum ivy_status status = IVY_OK;
	if (!current || current != unit->s_immediate) {
		status = switch_to_immediate_func(gen, unit, ENSURE_OPEN_BLOCK);
	}

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

	codegen_push_generator(gen, CODE_GENERATOR_RETURN, 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 unit_codegen_state *unit = (struct unit_codegen_state *)state;
	codegen_var_map_init(&unit->s_vars);

	return IVY_OK;
}

static enum ivy_status state_fini(
	struct ivy_codegen *gen, struct code_generator_state *state,
	struct code_generator_value *result)
{
	struct unit_codegen_state *unit = (struct unit_codegen_state *)state;
	codegen_var_map_fini(&unit->s_vars);

	struct mie_func *current = mie_builder_get_current_func(gen->c_builder);

	enum ivy_status status = IVY_OK;
	if (!current || current != unit->s_immediate) {
		status = switch_to_immediate_func(gen, unit, 0);
	}

	struct mie_block *block = mie_builder_get_current_block(gen->c_builder);
	if (!block->b_terminator) {
		mie_builder_ret(gen->c_builder, NULL);
	}

	return IVY_OK;
}

static enum ivy_status define_var(
	struct ivy_codegen *gen, struct code_generator_state *state,
	const char *ident, const struct codegen_var *var)
{
	struct unit_codegen_state *unit = (struct unit_codegen_state *)state;
	return codegen_var_map_put(&unit->s_vars, ident, var);
}

static enum ivy_status resolve_var(
	struct ivy_codegen *gen, struct code_generator_state *state,
	const char *ident, struct codegen_var *var)
{
	struct unit_codegen_state *unit = (struct unit_codegen_state *)state;
	struct codegen_var *result = NULL;
	enum ivy_status status
		= codegen_var_map_get(&unit->s_vars, ident, &result);

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

	memcpy(var, result, sizeof *var);
	return IVY_OK;
}

struct code_generator unit_generator = {
	.g_type = CODE_GENERATOR_UNIT,
	.g_state_size = sizeof(struct unit_codegen_state),
	.g_state_init = state_init,
	.g_state_fini = state_fini,
	.g_node_generators = {
		NODE_CODEGEN(RETURN, gen_return),
		NODE_CODEGEN(FOR_LOOP, gen_for_loop),
		NODE_CODEGEN(WHILE_LOOP, gen_while_loop),
		NODE_CODEGEN(UNIT_PACKAGE, gen_unit_package),
		NODE_CODEGEN(UNIT_IMPORT, gen_unit_import),
	},
	.g_expr_generator = gen_expr,
	.g_define_var = define_var,
	.g_resolve_var = resolve_var,
};