#include <mie/ctx.h>
#include <mie/diag/diag.h>
#include <mie/diag/highlight.h>
#include <mie/dialect/builtin.h>
#include <mie/dialect/dialect.h>
#include <mie/ir/block.h>
#include <mie/ir/op.h>
#include <mie/ir/region.h>
#include <mie/ir/register.h>

enum register_find_result {
	REG_FIND_NONE = 0,
	REG_FIND_USE_BEFORE_DEFINE,
	REG_FIND_UNDOMINATED,
	REG_FIND_ISOLATED,
};

bool mie_resolve_op_self(struct mie_op *op, struct mie_ctx *ctx)
{
	if (op->op_flags & MIE_OP_F_OP_RESOLVED) {
		return true;
	}

	const char *dialect_name = NULL, *op_name = NULL;

	char *dot = strchr(op->op_name, '.');
	if (dot) {
		*dot = 0;
		dialect_name = op->op_name;
		op_name = dot + 1;
	} else {
		dialect_name = NULL;
		op_name = op->op_name;
	}

	const struct mie_dialect *dialect = mie_ctx_get_dialect(ctx, dialect_name);
	if (dot) {
		*dot = '.';
	}

	/* dialect_name is no longer valid after this point */
	dialect_name = NULL;

	if (!dialect) {
		return false;
	}

	const struct mie_op_definition *op_info
		= mie_dialect_get_op(dialect, op_name);
	if (!op_info) {
		return false;
	}

	op->op_info = op_info;

	free(op->op_name);
	op->op_name = NULL;

	op->op_flags |= MIE_OP_F_OP_RESOLVED;

	return true;
}

static enum register_find_result find_register_wide(
	struct mie_op *op, const char *name, struct mie_register **out)
{
	struct mie_region *region = op->op_container->b_parent;
	enum register_find_result result = REG_FIND_UNDOMINATED;
	struct mie_register *reg = NULL;

	while (region) {
		reg = mie_region_find_register(region, name);
		if (reg) {
			*out = reg;
			break;
		}

		struct mie_op *op = region->r_parent;
		if (mie_op_has_trait(op, "builtin", "isolated-from-above")) {
			result = REG_FIND_ISOLATED;
		}

		region = op->op_container ? op->op_container->b_parent : NULL;
	}

	if (!reg) {
		return result;
	}

	if (reg->reg_block == op->op_container) {
		return REG_FIND_USE_BEFORE_DEFINE;
	}

	return result;
}

static bool resolve_arg(
	struct mie_op *op, struct mie_op_arg *arg, struct mie_ctx *ctx)
{
	const char *arg_name = arg->arg_unresolved.reg_name;
	struct mie_block *block = op->op_container;
	struct mie_op *search_start = op;
	struct mie_register *reg = NULL;
	bool cfg = (block->b_idom != NULL);

	while (block) {
		reg = mie_block_find_register(block, arg_name, search_start);
		if (reg) {
			break;
		}

		search_start = NULL;
		block = block->b_idom;
	}

	if (reg) {
		free(arg->arg_unresolved.reg_name);
		arg->arg_flags |= MIE_OP_F_ARG_RESOLVED;

		memset(&arg->arg_value, 0x0, sizeof arg->arg_value);
		arg->arg_value.u_reg = reg;
		arg->arg_value.u_user = op;
		b_queue_push_back(&reg->reg_use, &arg->arg_value.u_entry);
		return true;
	}

	enum register_find_result find_result
		= find_register_wide(op, arg_name, &reg);

	if (!cfg && REG_FIND_UNDOMINATED) {
		/* this isn't a cfg (yet), so ignore dominance-relaced resolution issues */
		free(arg->arg_unresolved.reg_name);
		arg->arg_flags |= MIE_OP_F_ARG_RESOLVED;

		memset(&arg->arg_value, 0x0, sizeof arg->arg_value);
		arg->arg_value.u_reg = reg;
		arg->arg_value.u_user = op;
		b_queue_push_back(&reg->reg_use, &arg->arg_value.u_entry);
		return true;
	}

	struct mie_diag *diag = mie_ctx_push_diag(
		ctx, op->op_src, &arg->arg_span.s_start, "builtin",
		MIE_BUILTIN_E_UNRESOLVED_VALUE);
	mie_diag_push_msg(diag, ctx, "builtin", MIE_BUILTIN_MSG_UNRESOLVED_VALUE);

	struct mie_diag_highlight hl[] = {
		{
			.hl_type = MIE_DIAG_HIGHLIGHT_ERROR,
			.hl_span = arg->arg_span,
		},
	};

	mie_diag_push_snippet(
		diag, arg->arg_span.s_start.c_row, arg->arg_span.s_end.c_row,
		NULL, 0, hl, 1);

	if (!reg) {
		return false;
	}

	struct mie_block *reg_container = reg->reg_block;

	switch (find_result) {
	case REG_FIND_ISOLATED:
		mie_diag_push_msg(
			diag, ctx, "builtin",
			MIE_BUILTIN_MSG_VALUE_DEFINED_OUTSIDE_ISOLATED_REGION);
		break;
	case REG_FIND_UNDOMINATED:
		mie_diag_push_msg(
			diag, ctx, "builtin",
			MIE_BUILTIN_MSG_VALUE_DEFINED_IN_NON_DOMINANT_BLOCK);
		break;
	case REG_FIND_USE_BEFORE_DEFINE:
		mie_diag_push_msg(
			diag, ctx, "builtin",
			MIE_BUILTIN_MSG_VALUE_DEFINED_AFTER_USE);
		break;
	default:
		return false;
	}

	hl[0].hl_type = MIE_DIAG_HIGHLIGHT_HINT;
	hl[0].hl_span = reg->reg_span;

	mie_diag_push_snippet(
		diag, reg->reg_span.s_start.c_row, reg->reg_span.s_end.c_row,
		NULL, 0, hl, 1);

	if (!reg_container) {
		return false;
	}

	hl[0].hl_span = reg_container->b_name.n_span;

	mie_diag_push_snippet(
		diag, reg_container->b_name.n_span.s_start.c_row,
		reg_container->b_name.n_span.s_end.c_row, NULL, 0, hl, 1);

	return false;
}

static bool resolve_successor(
	struct mie_op *op, struct mie_op_successor *s, struct mie_ctx *ctx)
{
	if (s->s_flags & MIE_OP_F_SUCCESSOR_RESOLVED) {
		return true;
	}

	struct mie_block *container = op->op_container;
	struct mie_region *region = container ? container->b_parent : NULL;
	if (!region) {
		return false;
	}

	struct mie_diag *diag = NULL;
	struct mie_block *dest = mie_region_find_block(region, s->s_block_name);
	if (!dest) {
		diag = mie_ctx_push_diag(
			ctx, op->op_src, &s->s_name_span.s_start, "builtin",
			MIE_BUILTIN_E_UNRESOLVED_SUCCESSOR);
		mie_diag_push_msg(
			diag, ctx, "builtin", MIE_BUILTIN_MSG_CANNOT_FIND_BLOCK);
		const struct mie_diag_highlight hl[] = {
			{
				.hl_type = MIE_DIAG_HIGHLIGHT_ERROR,
				.hl_span = s->s_name_span,
			},
		};
		mie_diag_push_snippet(
			diag, s->s_name_span.s_start.c_row,
			s->s_name_span.s_end.c_row, NULL, 0, hl, 1);
		return false;
	}

	free(s->s_block_name);
	s->s_block = dest;
	s->s_flags |= MIE_OP_F_SUCCESSOR_RESOLVED;

	return true;
}

bool mie_resolve_op_args(struct mie_op *op, struct mie_ctx *ctx)
{
	bool ok = true;

	for (size_t i = 0; i < MIE_VECTOR_COUNT(op->op_args); i++) {
		if (!resolve_arg(op, &op->op_args.items[i], ctx)) {
			ok = false;
		}
	}

	for (size_t i = 0; i < MIE_VECTOR_COUNT(op->op_successors); i++) {
		struct mie_op_successor *s = &op->op_successors.items[i];
		for (size_t k = 0; k < MIE_VECTOR_COUNT(s->s_args); k++) {
			if (!resolve_arg(op, &s->s_args.items[k], ctx)) {
				ok = false;
			}
		}
	}

	return ok;
}

bool mie_resolve_op_successors(struct mie_op *op, struct mie_ctx *ctx)
{
	bool ok = true;

	for (size_t i = 0; i < MIE_VECTOR_COUNT(op->op_successors); i++) {
		if (!resolve_successor(op, &op->op_successors.items[i], ctx)) {
			ok = false;
		}
	}

	return ok;
}