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

struct while_parser_state {
	struct parser_state s_base;
	bool s_inline;
	unsigned int s_prev_token;

	struct ivy_ast_node *s_body, *s_cond;

	struct ivy_ast_node *s_prev_node;
};

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

struct token_parse_result parse_while(struct ivy_parser *ctx, struct ivy_token *tok)
{
	struct while_parser_state *state
		= parser_get_state(ctx, struct while_parser_state);

	if (state->s_prev_node) {
		/* this is an inline while-loop */
		state->s_inline = true;
		state->s_body = state->s_prev_node;
		state->s_prev_node = NULL;
	} else {
		state->s_inline = false;
	}

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

	state->s_prev_token = IVY_KW_WHILE;
	expr->s_subexpr_depth = 1;
	expr->s_terminator = IVY_KW_DO;

	return PARSE_RESULT(IVY_OK, 0);
}

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

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

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

	/* previous component was the while-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 block. */
	struct block_parser_state *block
		= (struct block_parser_state *)parser_push_state(
			ctx, IVY_AST_BLOCK, 0);
	/* set the sub-expression depth to be non-zero so the expression parser doesn't consume the expression separator. */
	block->s_terminator = IVY_KW_END;

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

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

	if (state->s_prev_token != IVY_KW_DO) {
		/* expression can only follow else and then keywords. */
		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_KW_END;
	return PARSE_RESULT(IVY_OK, PARSE_REPEAT_TOKEN);
}

static enum ivy_status finalise_while_loop(struct while_parser_state *state)
{
	struct ivy_ast_while_loop_node *loop
		= (struct ivy_ast_while_loop_node *)state->s_base.s_node;

	if (state->s_inline) {
		/* we have just reached the end of either the while condition. */
		if (!state->s_prev_node) {
			/* while condition is empty. */
			return IVY_ERR_BAD_SYNTAX;
		}

		/* the condition and body are parsed in reverse order for inline loops. */
		loop->n_cond = state->s_prev_node;
		loop->n_body = state->s_body;
		state->s_prev_node = NULL;
		return IVY_OK;
	}

	if (!state->s_prev_node) {
		/* while body is empty. */
		return IVY_ERR_BAD_SYNTAX;
	}

	loop->n_cond = state->s_cond;
	loop->n_body = state->s_prev_node;

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

	return IVY_OK;
}

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

	if (!state->s_inline) {
		/* only inline while loop can be ended with punctuation. */
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

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

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

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

	if (state->s_inline) {
		/* inline while loop must be terminated with punctuation. */
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

	enum ivy_status status = finalise_while_loop(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 enum ivy_status add_child(
	struct parser_state *parent, struct ivy_ast_node *child)
{
	struct while_parser_state *state = (struct while_parser_state *)parent;

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

	state->s_prev_node = child;
	return IVY_OK;
}

static void collect_children(
	struct ivy_ast_node *node, struct ivy_ast_node_iterator *iterator)
{
	struct ivy_ast_while_loop_node *loop
		= (struct ivy_ast_while_loop_node *)node;

	if (loop->n_cond) {
		ast_node_iterator_enqueue_node(iterator, node, loop->n_cond);
	}

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

struct ast_node_type while_loop_node_ops = {
	.n_init_state = init_state,
	.n_add_child = add_child,
	.n_collect_children = collect_children,
	.n_state_size = sizeof(struct while_parser_state),
	.n_node_size = sizeof(struct ivy_ast_while_loop_node),
	.n_keyword_parsers = {
		KW_PARSER(WHILE, parse_while),
		KW_PARSER(DO, parse_do),
		KW_PARSER(END, parse_end),
	},
	.n_symbol_parsers = {
		SYM_PARSER(DOT, parse_punct_terminator),
		SYM_PARSER(COMMA, parse_punct_terminator),
		SYM_PARSER(RIGHT_PAREN, parse_punct_terminator),
		SYM_PARSER(BANG, parse_punct_terminator),
	},
	.n_expr_parser = {
		.expr_begin = parse_expr_begin,
	},
};