#include "codegen.h"

#include <mie/arg.h>
#include <mie/block.h>
#include <mie/func.h>
#include <mie/module.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

enum lambda_part {
	LAMBDA_NONE,
	LAMBDA_START,
	LAMBDA_ARG,
	LAMBDA_BLOCK,
};

struct lambda_codegen_state {
	struct code_generator_state s_base;
	enum lambda_part s_prev_part;
	struct mie_func *s_func;
	struct mie_value *s_ctx_internal;
	struct mie_value *s_ctx_external;
	struct mie_block *s_outer_block;
	struct codegen_var_map s_args;
	struct codegen_var_map s_captured_vars;
};

static enum ivy_status switch_to_lambda_func(
	struct ivy_codegen *gen, struct lambda_codegen_state *lambda)
{
	if (!lambda->s_outer_block) {
		lambda->s_outer_block
			= mie_builder_get_current_block(gen->c_builder);
	}

	struct mie_block *block = mie_func_get_last_block(lambda->s_func);
	if (!block) {
		block = mie_func_create_block(lambda->s_func, "entry");
		mie_func_insert_block(lambda->s_func, block, NULL);
	}

	if (mie_block_is_terminated(block)) {
		block = mie_func_create_block(lambda->s_func, NULL);
		mie_func_insert_block(lambda->s_func, block, NULL);
	}

	mie_builder_set_insert_point(gen->c_builder, block);

	return IVY_OK;
}

static enum ivy_status switch_to_outer_block(
	struct ivy_codegen *gen, struct lambda_codegen_state *lambda)
{
	mie_builder_set_insert_point(gen->c_builder, lambda->s_outer_block);
	return IVY_OK;
}

static struct code_generator_result gen_lambda(
	struct ivy_codegen *gen, struct code_generator_state *state,
	struct ivy_ast_node *node, size_t depth)
{
	struct lambda_codegen_state *lambda = (struct lambda_codegen_state *)state;
	if (lambda->s_prev_part != LAMBDA_NONE) {
		return CODEGEN_RESULT_ERR(IVY_ERR_BAD_SYNTAX);
	}

	struct mie_type *id = mie_ctx_get_type(gen->c_ctx, MIE_TYPE_ID);
	struct mie_type *func_type = mie_ctx_get_type(gen->c_ctx, MIE_TYPE_FUNC);

	lambda->s_prev_part = LAMBDA_START;
	lambda->s_func = mie_func_create(MIE_FUNC_STATIC, id);
	lambda->s_ctx_internal = mie_func_add_arg(lambda->s_func, id, "ctx");

	struct mie_value *lambda_class
		= mie_builder_get_data_ptr(gen->c_builder, "Lambda");
	struct mie_value *sel = mie_ctx_get_selector(gen->c_ctx, "_M3new4funcE");
	struct mie_value *func_value = mie_builder_load(
		gen->c_builder, id, MIE_VALUE(lambda->s_func), NULL);
	struct mie_value *args[] = {
		func_value,
	};

	lambda->s_ctx_external = mie_builder_msg(
		gen->c_builder, id, lambda_class, sel, args,
		sizeof args / sizeof args[0], 0, NULL);

	switch_to_lambda_func(gen, lambda);

	return CODEGEN_RESULT_OK(0);
}

static struct code_generator_result gen_ident(
	struct ivy_codegen *gen, struct code_generator_state *state,
	struct ivy_ast_node *node, size_t depth)
{
	struct lambda_codegen_state *lambda = (struct lambda_codegen_state *)state;
	if (lambda->s_prev_part != LAMBDA_START
	    && lambda->s_prev_part != LAMBDA_ARG) {
		return CODEGEN_RESULT_ERR(IVY_ERR_BAD_SYNTAX);
	}

	lambda->s_prev_part = LAMBDA_ARG;

	struct ivy_ast_ident_node *ident = (struct ivy_ast_ident_node *)node;
	const char *arg_name = ident->n_content->t_str;
	struct mie_type *id = mie_ctx_get_type(gen->c_ctx, MIE_TYPE_ID);
	struct mie_value *arg = mie_func_add_arg(lambda->s_func, id, arg_name);

	struct codegen_var arg_var = {
		.v_node = node,
		.v_type = id,
		.v_value = arg,
	};

	codegen_var_map_put(&lambda->s_args, arg_name, &arg_var);

	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 lambda_codegen_state *lambda = (struct lambda_codegen_state *)state;
	if (lambda->s_prev_part != LAMBDA_START
	    && lambda->s_prev_part != LAMBDA_ARG) {
		return CODEGEN_RESULT_ERR(IVY_ERR_BAD_SYNTAX);
	}

	codegen_push_generator(gen, CODE_GENERATOR_BLOCK, 0, NULL);

	return CODEGEN_RESULT_OK(0);
}

static enum ivy_status state_init(
	struct ivy_codegen *gen, struct code_generator_state *state,
	uintptr_t argv, void *argp)
{
	struct lambda_codegen_state *lambda = (struct lambda_codegen_state *)state;
	codegen_var_map_init(&lambda->s_args);
	codegen_var_map_init(&lambda->s_captured_vars);
	lambda->s_prev_part = LAMBDA_NONE;
	return IVY_OK;
}

static enum ivy_status state_fini(
	struct ivy_codegen *gen, struct code_generator_state *state,
	struct mie_value **result)
{
	struct lambda_codegen_state *lambda = (struct lambda_codegen_state *)state;

	switch_to_outer_block(gen, lambda);
	mie_module_add_function(gen->c_module, lambda->s_func, ".anon");
	*result = MIE_VALUE(lambda->s_ctx_external);

	codegen_var_map_fini(&lambda->s_args);
	codegen_var_map_fini(&lambda->s_captured_vars);

	return IVY_OK;
}

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

	if (status == IVY_OK) {
		memcpy(var, result, sizeof *var);
		return status;
	}

	status = codegen_var_map_get(&lambda->s_captured_vars, ident, &result);

	if (status == IVY_OK) {
		memcpy(var, result, sizeof *var);
		return status;
	}

	return IVY_ERR_NO_ENTRY;
}

static enum ivy_status capture_var(
	struct ivy_codegen *gen, struct code_generator_state *state,
	const char *ident, struct codegen_var *var)
{
	struct lambda_codegen_state *lambda = (struct lambda_codegen_state *)state;

	struct mie_type *str_type = mie_ctx_get_type(gen->c_ctx, MIE_TYPE_STR);
	struct mie_type *void_type = mie_ctx_get_type(gen->c_ctx, MIE_TYPE_VOID);
	struct mie_value *ident_ptr
		= mie_builder_get_string_ptr(gen->c_builder, ident);

	switch_to_outer_block(gen, lambda);

	struct mie_value *ident_str
		= mie_builder_load(gen->c_builder, str_type, ident_ptr, NULL);

	struct mie_value *capture_value = codegen_load_variable(gen, var);
	struct mie_value *put_sel
		= mie_ctx_get_selector(gen->c_ctx, "_M02at3putE");
	struct mie_value *put_args[] = {
		ident_str,
		capture_value,
	};

	struct mie_value *capture_store = mie_builder_msg(
		gen->c_builder, void_type, lambda->s_ctx_external, put_sel,
		put_args, sizeof put_args / sizeof put_args[0],
		MIE_BUILDER_IGNORE_RESULT, NULL);

	switch_to_lambda_func(gen, lambda);

	struct codegen_var capture = {
		.v_type = var->v_type,
		.v_node = var->v_node,
		.v_flags = CODEGEN_VAR_F_CAPTURE | CODEGEN_VAR_F_PTR,
	};

	ident_str = mie_builder_load(gen->c_builder, str_type, ident_ptr, NULL);

	struct mie_value *at_sel = mie_ctx_get_selector(gen->c_ctx, "_M02atE");
	struct mie_value *at_args[] = {
		ident_str,
	};

	struct mie_value *stack_ptr
		= mie_builder_alloca(gen->c_builder, var->v_type, ident);

	struct mie_value *value = mie_builder_msg(
		gen->c_builder, var->v_type, lambda->s_ctx_internal, at_sel,
		at_args, sizeof at_args / sizeof at_args[0], 0, NULL);
	mie_builder_store(gen->c_builder, value, stack_ptr);

	capture.v_value = stack_ptr;

	*var = capture;

	return IVY_OK;
}

static struct code_generator_result value_received(
	struct ivy_codegen *gen, struct code_generator_state *state,
	struct mie_value *value)
{
	return CODEGEN_RESULT_OK(0);
}

struct code_generator lambda_generator = {
	.g_type = CODE_GENERATOR_LAMBDA,
	.g_state_size = sizeof(struct lambda_codegen_state),
	.g_state_init = state_init,
	.g_state_fini = state_fini,
	.g_value_received = value_received,
	.g_resolve_var = resolve_var,
	.g_capture_var = capture_var,
	.g_node_generators = {
		NODE_CODEGEN(LAMBDA, gen_lambda),
		NODE_CODEGEN(IDENT, gen_ident),
		NODE_CODEGEN(BLOCK, gen_block),
	},
};