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

struct try_parser_state {
	struct parser_state s_base;
	unsigned int s_prev_token;

	struct ivy_ast_node *s_prev_node;

	struct ivy_ast_block_node *s_try, *s_finally;

	b_queue s_catch_branches;
	struct ivy_ast_try_catch_node *s_cur_catch_branch;
};

static enum ivy_status flush_catch_branch(struct try_parser_state *state)
{
	b_queue_push_back(
		&state->s_catch_branches,
		&state->s_cur_catch_branch->n_base.n_entry);

	state->s_cur_catch_branch
		= (struct ivy_ast_try_catch_node *)ast_node_create(
			IVY_AST_TRY_CATCH);

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

	return IVY_OK;
}

static enum ivy_status finalise_try(struct try_parser_state *state)
{
	struct ivy_ast_try_node *try
		= (struct ivy_ast_try_node *)state->s_base.s_node;
	try->n_try = (struct ivy_ast_node *)state->s_try;
	try->n_finally = (struct ivy_ast_node *)state->s_finally;
	try->n_catch = state->s_catch_branches;

	state->s_try = NULL;
	state->s_finally = NULL;
	state->s_catch_branches = B_QUEUE_INIT;

	return IVY_OK;
}

struct token_parse_result parse_catch(struct ivy_parser *ctx, struct ivy_token *tok)
{
	struct try_parser_state *try
		= parser_get_state(ctx, struct try_parser_state);

	if (!try->s_prev_node || try->s_prev_node->n_type != IVY_AST_BLOCK) {
		/* there must always be a block between a catch and the keyword
		 * that came before it */
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

	enum ivy_status status = IVY_OK;

	if (try->s_prev_token == IVY_KW_TRY) {
		try->s_try = (struct ivy_ast_block_node *)try->s_prev_node;
		try->s_prev_node = NULL;
	} else if (try->s_prev_token == IVY_KW_IN) {
		try->s_cur_catch_branch->n_block = try->s_prev_node;
		try->s_prev_node = NULL;

		status = flush_catch_branch(try);

		if (status != IVY_OK) {
			return PARSE_RESULT(status, 0);
		}
	} else {
		/* catch can only come after `try` or `in` */
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

	/* next input should be the catch pattern */
	struct expr_parser_state *pattern
		= (struct expr_parser_state *)parser_push_state(
			ctx, IVY_AST_EXPR, 0);

	expr_add_terminator(pattern, IVY_KW_IN);

	try->s_prev_token = IVY_KW_CATCH;

	return PARSE_RESULT(IVY_OK, 0);
}

static struct token_parse_result parse_in(
	struct ivy_parser *ctx, struct ivy_token *tok)
{
	struct try_parser_state *try
		= parser_get_state(ctx, struct try_parser_state);

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

	if (!try->s_prev_node) {
		/* there must always be an expression between catch and in */
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

	try->s_cur_catch_branch->n_pattern = try->s_prev_node;
	try->s_prev_node = NULL;
	try->s_prev_token = IVY_KW_IN;

	struct block_parser_state *block
		= (struct block_parser_state *)parser_push_state(
			ctx, IVY_AST_BLOCK, 0);
	block_add_terminator(block, IVY_KW_CATCH);
	block_add_terminator(block, IVY_KW_FINALLY);
	block_add_terminator(block, IVY_KW_END);

	return PARSE_RESULT(IVY_OK, 0);
}

struct token_parse_result parse_finally(
	struct ivy_parser *ctx, struct ivy_token *tok)
{
	struct try_parser_state *try
		= parser_get_state(ctx, struct try_parser_state);

	if (try->s_prev_token != IVY_KW_IN && try->s_prev_token != IVY_KW_TRY) {
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

	if (!try->s_prev_node || try->s_prev_node->n_type != IVY_AST_BLOCK) {
		/* there must always be a block between a catch and the keyword
		 * that came before it */
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

	enum ivy_status status = IVY_OK;

	if (try->s_prev_token == IVY_KW_TRY) {
		try->s_try = (struct ivy_ast_block_node *)try->s_prev_node;
		try->s_prev_node = NULL;
	} else if (try->s_prev_token == IVY_KW_IN) {
		try->s_cur_catch_branch->n_block = try->s_prev_node;
		try->s_prev_node = NULL;

		status = flush_catch_branch(try);

		if (status != IVY_OK) {
			return PARSE_RESULT(status, 0);
		}
	} else {
		/* finally can only come after `try` or `in` */
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

	try->s_prev_token = IVY_KW_FINALLY;

	struct block_parser_state *block
		= (struct block_parser_state *)parser_push_state(
			ctx, IVY_AST_BLOCK, 0);
	block_add_terminator(block, IVY_KW_END);

	return PARSE_RESULT(IVY_OK, 0);
}

struct token_parse_result parse_end(struct ivy_parser *ctx, struct ivy_token *tok)
{
	struct try_parser_state *try
		= parser_get_state(ctx, struct try_parser_state);

	if (!try->s_prev_node || try->s_prev_node->n_type != IVY_AST_BLOCK) {
		/* there must always be a block between an end and the keyword
		 * that came before it */
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

	enum ivy_status status = IVY_OK;

	if (try->s_prev_token == IVY_KW_TRY) {
		try->s_try = (struct ivy_ast_block_node *)try->s_prev_node;
		try->s_prev_node = NULL;
	} else if (try->s_prev_token == IVY_KW_FINALLY) {
		try->s_finally = (struct ivy_ast_block_node *)try->s_prev_node;
		try->s_prev_node = NULL;
	} else if (try->s_prev_token == IVY_KW_IN) {
		try->s_cur_catch_branch->n_block = try->s_prev_node;
		status = flush_catch_branch(try);

		if (status != IVY_OK) {
			return PARSE_RESULT(status, 0);
		}
	} else {
		/* catch can only come after `try`, `in`, or `finally` */
		return PARSE_RESULT(IVY_ERR_BAD_SYNTAX, 0);
	}

	finalise_try(try);
	parser_pop_state(ctx, STATE_ADD_NODE_TO_PARENT);

	return PARSE_RESULT(IVY_OK, 0);
}

static void init_state(struct ivy_parser *ctx, struct parser_state *sp, uintptr_t arg)
{
	struct try_parser_state *state = (struct try_parser_state *)sp;
	state->s_prev_token = IVY_KW_TRY;

	state->s_cur_catch_branch
		= (struct ivy_ast_try_catch_node *)ast_node_create(
			IVY_AST_TRY_CATCH);

	struct block_parser_state *block
		= (struct block_parser_state *)parser_push_state(
			ctx, IVY_AST_BLOCK, 0);
	block_add_terminator(block, IVY_KW_CATCH);
	block_add_terminator(block, IVY_KW_FINALLY);
	block_add_terminator(block, IVY_KW_END);
}

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

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

	state->s_prev_node = child;
	return IVY_OK;
}

static void try_collect_children(
	struct ivy_ast_node *node, struct ivy_ast_node_iterator *iterator)
{
	struct ivy_ast_try_node *try = (struct ivy_ast_try_node *)node;

	ast_node_iterator_enqueue_node(iterator, node, try->n_try);

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

	if (try->n_finally) {
		ast_node_iterator_enqueue_node(iterator, node, try->n_finally);
	}
}

static void try_catch_collect_children(
	struct ivy_ast_node *node, struct ivy_ast_node_iterator *iterator)
{
	struct ivy_ast_try_catch_node *catch
		= (struct ivy_ast_try_catch_node *)node;

	ast_node_iterator_enqueue_node(iterator, node, catch->n_pattern);
	ast_node_iterator_enqueue_node(iterator, node, catch->n_block);
}

struct ast_node_type try_node_ops = {
	.n_init_state = init_state,
	.n_add_child = add_child,
	.n_collect_children = try_collect_children,
	.n_state_size = sizeof(struct try_parser_state),
	.n_node_size = sizeof(struct ivy_ast_try_node),
	.n_keyword_parsers = {
		KW_PARSER(CATCH, parse_catch),
		KW_PARSER(IN, parse_in),
		KW_PARSER(FINALLY, parse_finally),
		KW_PARSER(END, parse_end),
	},
};

struct ast_node_type try_catch_node_ops = {
	.n_state_size = sizeof(struct parser_state),
	.n_node_size = sizeof(struct ivy_ast_try_catch_node),
	.n_collect_children = try_catch_collect_children,
};