#include "block.h"
#include "expr/expr.h"
#include "iterate.h"

struct match_parser_state {
	struct parser_state s_base;
	unsigned int s_prev_token;

	struct ivy_ast_node *s_cond;

	struct ivy_ast_cond_node *s_cur_branch;
	b_queue s_branches;

	struct ivy_ast_node *s_prev_node;
};

static void init_state(struct ivy_parser *ctx, struct parser_state *sp, uintptr_t arg)
{
	struct match_parser_state *state = (struct match_parser_state *)sp;
	state->s_prev_node = (struct ivy_ast_node *)arg;
}

struct token_parse_result parse_match(struct ivy_parser *ctx, struct ivy_token *tok)
{
	struct match_parser_state *state
		= parser_get_state(ctx, struct match_parser_state);

	if (state->s_cur_branch || state->s_prev_node) {
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

	state->s_cur_branch
		= (struct ivy_ast_cond_node *)ast_node_create(IVY_AST_COND);
	if (!state->s_cur_branch) {
		return PARSE_RESULT(IVY_ERR_NO_MEMORY, 0);
	}

	struct expr_parser_state *expr
		= (struct expr_parser_state *)parser_push_state(
			ctx, IVY_AST_EXPR, 0);

	state->s_prev_token = IVY_KW_MATCH;
	expr->s_subexpr_depth = 1;
	expr_add_terminator(expr, IVY_KW_IN);

	return PARSE_RESULT(IVY_OK, 0);
}

static enum ivy_status flush_current_branch(struct match_parser_state *state)
{
	if (!state->s_cur_branch) {
		return IVY_ERR_INTERNAL_FAILURE;
	}

	b_queue_push_back(&state->s_branches, &state->s_cur_branch->n_base.n_entry);

	state->s_cur_branch
		= (struct ivy_ast_cond_node *)ast_node_create(IVY_AST_COND);

	if (!state->s_cur_branch) {
		return IVY_ERR_NO_MEMORY;
	}

	return IVY_OK;
}

struct token_parse_result parse_in(struct ivy_parser *ctx, struct ivy_token *tok)
{
	struct match_parser_state *state
		= parser_get_state(ctx, struct match_parser_state);

	if (!state->s_cur_branch) {
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

	if (state->s_prev_token != IVY_KW_MATCH) {
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

	/* previous component was the match-condition. */
	if (!state->s_prev_node) {
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

	state->s_cond = state->s_prev_node;
	state->s_prev_node = NULL;

	/* next component will be a branch condition. */
	struct block_parser_state *block
		= (struct block_parser_state *)parser_push_state(
			ctx, IVY_AST_EXPR, 0);

	state->s_prev_token = IVY_KW_IN;
	return PARSE_RESULT(IVY_OK, 0);
}

static enum ivy_status finalise_match(struct match_parser_state *state)
{
	struct ivy_ast_match_node *match
		= (struct ivy_ast_match_node *)state->s_base.s_node;

	/* we have just reached the 'end' keyword. */
	if (!state->s_cur_branch) {
		/* not currently parsing a conditional branch. */
		return IVY_ERR_BAD_SYNTAX;
	}

	match->n_cond = state->s_cond;
	match->n_branches = state->s_branches;
	state->s_branches = B_QUEUE_INIT;
	return IVY_OK;
}

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

	if (state->s_prev_token != IVY_KW_IN
	    && state->s_prev_token != IVY_SYM_COMMA) {
		/* this token can only appear after the `in` keyword and an expression. */
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

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

	state->s_cur_branch->n_cond = state->s_prev_node;
	state->s_prev_node = NULL;
	state->s_prev_token = IVY_SYM_EQUAL_RIGHT_ANGLE;

	/* the next component is the branch body */
#if 0
	struct expr_parser_state *expr
		= (struct expr_parser_state *)parser_push_state(
			ctx, IVY_AST_EXPR, 0);

	expr_add_terminator(expr, IVY_KW_END);
	expr_add_terminator(expr, IVY_SYM_COMMA);
	expr->s_subexpr_depth = 1;
#endif
	return PARSE_RESULT(IVY_OK, 0);
}

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

	if (state->s_prev_token != IVY_SYM_EQUAL_RIGHT_ANGLE) {
		/* this token can only appear after the `=>` symbol and an expression. */
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

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

	state->s_cur_branch->n_body = state->s_prev_node;
	flush_current_branch(state);
	state->s_prev_node = NULL;
	state->s_prev_token = IVY_SYM_COMMA;

	/* the next component is a branch condition */
	struct expr_parser_state *expr
		= (struct expr_parser_state *)parser_push_state(
			ctx, IVY_AST_EXPR, 0);

	expr->s_subexpr_depth = 1;
	return PARSE_RESULT(IVY_OK, 0);
}

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

	/* end can only be used after the '=>' symbol and an expression */
	if (state->s_prev_token != IVY_SYM_EQUAL_RIGHT_ANGLE) {
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

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

	state->s_cur_branch->n_body = state->s_prev_node;
	flush_current_branch(state);
	state->s_prev_node = NULL;
	state->s_prev_token = IVY_KW_END;

	enum ivy_status status = finalise_match(state);
	if (status != IVY_OK) {
		return PARSE_RESULT(status, 0);
	}

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

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

	if (state->s_prev_token != IVY_SYM_EQUAL_RIGHT_ANGLE) {
		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);

	return PARSE_RESULT(IVY_OK, 0);
}

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

	if (state->s_prev_token != IVY_SYM_EQUAL_RIGHT_ANGLE) {
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

	struct expr_parser_state *expr
		= (struct expr_parser_state *)parser_push_state(
			ctx, IVY_AST_EXPR, 0);

	if (tok->t_type != IVY_TOK_KEYWORD) {
		expr_add_terminator(expr, IVY_KW_END);
	}

	expr_add_terminator(expr, IVY_SYM_COMMA);
	expr->s_subexpr_depth = 1;
	return PARSE_RESULT(IVY_OK, PARSE_REPEAT_TOKEN);
}

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

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

	state->s_prev_node = child;
	return IVY_OK;
}

static void match_collect_children(
	struct ivy_ast_node *node, struct ivy_ast_node_iterator *iterator)
{
	struct ivy_ast_match_node *match = (struct ivy_ast_match_node *)node;

	ast_node_iterator_enqueue_node(iterator, node, match->n_cond);

	b_queue_iterator it = {0};
	b_queue_foreach (&it, &match->n_branches) {
		struct ivy_ast_node *branch
			= b_unbox(struct ivy_ast_node, it.entry, n_entry);
		ast_node_iterator_enqueue_node(iterator, node, branch);
	}
}

struct ast_node_type match_node_ops = {
	.n_init_state = init_state,
	.n_add_child = add_child,
	.n_collect_children = match_collect_children,
	.n_state_size = sizeof(struct match_parser_state),
	.n_node_size = sizeof(struct ivy_ast_match_node),
	.n_expr_parser = parse_expr,
	.n_keyword_parsers = {
		KW_PARSER(MATCH, parse_match),
		KW_PARSER(IN, parse_in),
		KW_PARSER(END, parse_end),
		KW_PARSER(DO, parse_block),
	},
	.n_symbol_parsers = {
		SYM_PARSER(EQUAL_RIGHT_ANGLE, parse_arrow),
		SYM_PARSER(COMMA, parse_comma),
	},
};