#include "block.h"
#include "ctx.h"
#include "iterate.h"
#include "node.h"

#include <ivy/lang/lex.h>

struct lambda_parser_state {
	struct parser_state s_base;

	struct ivy_ast_node *s_prev_node;
	unsigned int s_prev;

	b_queue s_args;
	struct ivy_ast_node *s_body;
};

static struct token_parse_result parse_expr_begin(
	struct ivy_parser *ctx, struct ivy_token *tok)
{
	struct lambda_parser_state *state
		= parser_get_state(ctx, struct lambda_parser_state);

	if (state->s_prev != IVY_SYM_LEFT_BRACKET && state->s_prev != IVY_SYM_PIPE) {
		/* unexpected expression. */
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

	struct block_parser_state *block
		= (struct block_parser_state *)parser_push_state(
			ctx, IVY_AST_BLOCK, 0);
	block->s_terminator = IVY_SYM_RIGHT_BRACKET;
	return PARSE_RESULT(IVY_OK, PARSE_REPEAT_TOKEN);
}

static struct token_parse_result parse_ident(
	struct ivy_parser *ctx, struct ivy_token *tok)
{
	struct lambda_parser_state *state
		= parser_get_state(ctx, struct lambda_parser_state);

	if (state->s_prev == IVY_SYM_LEFT_BRACKET || state->s_prev == IVY_SYM_PIPE) {
		return parse_expr_begin(ctx, tok);
	}

	if (state->s_prev != IVY_SYM_COLON) {
		/* the only idents we're expecting are parameter names,
		 * which are always preceded with a colon */
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

	struct ivy_ast_ident_node *arg
		= (struct ivy_ast_ident_node *)ast_node_create(IVY_AST_IDENT);
	if (!arg) {
		return PARSE_RESULT(IVY_ERR_NO_MEMORY, 0);
	}

	arg->n_content = tok;
	b_queue_push_back(&state->s_args, &arg->n_base.n_entry);
	state->s_prev = IVY_TOK_IDENT;

	return PARSE_RESULT(IVY_OK, 0);
}

static struct token_parse_result parse_colon(
	struct ivy_parser *ctx, struct ivy_token *tok)
{
	struct lambda_parser_state *state
		= parser_get_state(ctx, struct lambda_parser_state);

	if (state->s_prev != IVY_SYM_LEFT_BRACKET
	    && state->s_prev != IVY_TOK_IDENT) {
		/* colons can only be used before arg names */
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

	state->s_prev = IVY_SYM_COLON;
	return PARSE_RESULT(IVY_OK, 0);
}

static struct token_parse_result parse_pipe(
	struct ivy_parser *ctx, struct ivy_token *tok)
{
	struct lambda_parser_state *state
		= parser_get_state(ctx, struct lambda_parser_state);

	if (state->s_prev != IVY_TOK_IDENT) {
		/* pipe can only be used after an arg name. */
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

	state->s_prev = IVY_SYM_PIPE;
	return PARSE_RESULT(IVY_OK, 0);
}

static void finalise_lambda(struct lambda_parser_state *state)
{
	struct ivy_ast_lambda_node *lambda
		= (struct ivy_ast_lambda_node *)state->s_base.s_node;

	lambda->n_arg = state->s_args;
	lambda->n_body = state->s_body;

	state->s_args = B_QUEUE_INIT;
	state->s_body = NULL;
}

static struct token_parse_result parse_right_bracket(
	struct ivy_parser *ctx, struct ivy_token *tok)
{
	struct lambda_parser_state *state
		= parser_get_state(ctx, struct lambda_parser_state);

	if (!state->s_prev_node) {
		/* empty lambda */
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

	state->s_body = state->s_prev_node;
	state->s_prev_node = NULL;

	finalise_lambda(state);
	parser_pop_state(ctx, STATE_ADD_NODE_TO_PARENT);
	return PARSE_RESULT(IVY_OK, 0);
}

static enum ivy_status add_child(
	struct parser_state *parent, struct ivy_ast_node *child)
{
	struct lambda_parser_state *state = (struct lambda_parser_state *)parent;

	if (state->s_prev_node) {
		return IVY_ERR_BAD_SYNTAX;
	}

	state->s_prev_node = child;
	return IVY_OK;
}

static void init_state(struct ivy_parser *ctx, struct parser_state *sp, uintptr_t arg)
{
	struct lambda_parser_state *state = (struct lambda_parser_state *)sp;
	state->s_prev = IVY_SYM_LEFT_BRACKET;
}

static void collect_children(
	struct ivy_ast_node *node, struct ivy_ast_node_iterator *iterator)
{
	struct ivy_ast_lambda_node *lambda = (struct ivy_ast_lambda_node *)node;
	b_queue_iterator it = {0};
	b_queue_foreach (&it, &lambda->n_arg) {
		struct ivy_ast_node *expr
			= b_unbox(struct ivy_ast_node, it.entry, n_entry);
		ast_node_iterator_enqueue_node(iterator, node, expr);
	}

	if (lambda->n_body) {
		ast_node_iterator_enqueue_node(iterator, node, lambda->n_body);
	}
}

struct ast_node_type lambda_node_ops = {
	.n_add_child = add_child,
	.n_init_state = init_state,
	.n_collect_children = collect_children,
	.n_state_size = sizeof(struct lambda_parser_state),
	.n_node_size = sizeof(struct ivy_ast_lambda_node),
	.n_token_parsers = {
		TOK_PARSER(IDENT, parse_ident),
	},
	.n_symbol_parsers = {
		SYM_PARSER(RIGHT_BRACKET, parse_right_bracket),
		SYM_PARSER(PIPE, parse_pipe),
		SYM_PARSER(COLON, parse_colon),
	},
	.n_expr_parser = {
		.expr_begin = parse_expr_begin,
	},
};