#include "assembler.h"

#include <blue/core/hash.h>
#include <blue/ds/hashmap.h>
#include <blue/ds/number.h>
#include <blue/ds/string.h>
#include <ivy/asm/bin.h>
#include <ivy/asm/instr.h>
#include <ivy/asm/lex.h>
#include <ivy/ident.h>
#include <ivy/selector.h>
#include <stdlib.h>
#include <string.h>

enum label_state {
	LABEL_UNRESOLVED = 0,
	LABEL_RESOLVED,
};

struct label {
	enum label_state l_state;
	b_queue_entry l_entry;
	const struct ivy_asm_token *l_name;
	size_t l_offset;
};

struct block_assembler_scope {
	struct assembler_scope s_base;
	b_hashmap *s_labels;
	b_queue s_label_refs;
	size_t s_text_start;
};

static struct label *label_create(enum label_state state)
{
	struct label *out = malloc(sizeof *out);
	if (!out) {
		return NULL;
	}

	memset(out, 0x0, sizeof *out);

	out->l_state = state;

	return out;
}

static void label_destroy(struct label *label)
{
	free(label);
}

static enum ivy_status init_scope(
	struct ivy_assembler *as, struct assembler_scope *scope,
	const ivy_assembler_attrib_table attrib)
{
	struct ivy_bin_block header = {0};

	struct block_assembler_scope *c = (struct block_assembler_scope *)scope;
	c->s_labels = b_hashmap_create(NULL, NULL);
	c->s_text_start = ivy_assembler_get_ptr(as);

	header.b_index = b_i32_htob((uint32_t)attrib[IVY_ASM_ATTRIB_INDEX]);
	assembler_write_data(as, &header, sizeof header);

	return IVY_OK;
}

static struct label *get_label(struct block_assembler_scope *scope, const char *name)
{
	b_hashmap_key key = {
		.key_data = name,
		.key_size = strlen(name),
	};

	const b_hashmap_value *v = b_hashmap_get(scope->s_labels, &key);
	if (!v) {
		return NULL;
	}

	return v->value_data;
}

static enum ivy_status resolve_label_refs(
	struct ivy_assembler *as, struct assembler_scope *scope)
{
	struct block_assembler_scope *c = (struct block_assembler_scope *)scope;
	enum ivy_status status = IVY_OK;
	size_t nr_read = 0;

	b_queue_entry *entry = b_queue_first(&c->s_label_refs);
	while (entry) {
		struct label *label_ref = b_unbox(struct label, entry, l_entry);
		struct label *label_dest = get_label(c, label_ref->l_name->t_str);

		if (!label_dest) {
			printf("cannot resolve label '%s'\n",
			       label_ref->l_name->t_str);
			return IVY_ERR_NO_ENTRY;
		}

		b_i32 x;
		status = assembler_read_data_at(
			as, &x, label_ref->l_offset, sizeof x, &nr_read);

		if (status != IVY_OK || nr_read != sizeof x) {
			printf("cannot fix label reference at 0x%zx\n",
			       label_ref->l_offset);
			return IVY_ERR_IO_FAILURE;
		}

		uint32_t instr = b_i32_btoh(x);
		instr = R_SET_D3(instr, label_dest->l_offset);
		x = b_i32_htob(instr);

		assembler_write_data_at(as, &x, label_ref->l_offset, sizeof x);
		entry = b_queue_next(entry);
	}

	return status;
}

static enum ivy_status finish_scope(
	struct ivy_assembler *as, struct assembler_scope *scope)
{
	return resolve_label_refs(as, scope);
}

static enum ivy_status put_instr(
	struct ivy_assembler *as, struct assembler_scope *scope,
	const struct ivy_instr *instr)
{
	uint32_t v = 0;
	v = R_SET_OPCODE(v, instr->i_op->i_opcode);

	switch (instr->i_op->i_layout) {
	case IVY_INSTR_X:
		break;
	case IVY_INSTR_D:
		v |= R_SET_D3(v, instr->i_arg[0]);
		break;
	case IVY_INSTR_R1:
		v |= R_SET_R1(v, instr->i_arg[0]);
		break;
	case IVY_INSTR_R1D:
		v |= R_SET_R1(v, instr->i_arg[0]);
		v |= R_SET_D2(v, instr->i_arg[1]);
		break;
	case IVY_INSTR_R2:
		v |= R_SET_R1(v, instr->i_arg[0]);
		v |= R_SET_R2(v, instr->i_arg[1]);
		break;
	case IVY_INSTR_R2D:
		v |= R_SET_R1(v, instr->i_arg[0]);
		v |= R_SET_R2(v, instr->i_arg[1]);
		v |= R_SET_D1(v, instr->i_arg[2]);
		break;
	case IVY_INSTR_R3:
		v |= R_SET_R1(v, instr->i_arg[0]);
		v |= R_SET_R2(v, instr->i_arg[1]);
		v |= R_SET_R3(v, instr->i_arg[2]);
		break;
	default:
		return IVY_ERR_INTERNAL_FAILURE;
	}

	b_i32 x = b_i32_htob(v);
	assembler_write_data(as, &x, sizeof x);
	return IVY_OK;
}

static enum ivy_status put_label(
	struct ivy_assembler *as, struct assembler_scope *scope,
	const struct ivy_asm_token *label_name, size_t label_offset)
{
	struct block_assembler_scope *c = (struct block_assembler_scope *)scope;
	struct label *label = label_create(LABEL_RESOLVED);

	label->l_name = label_name;
	label->l_offset = label_offset - c->s_text_start;

	b_hashmap_key key = {
		.key_data = label_name->t_str,
		.key_size = strlen(label_name->t_str),
	};

	b_hashmap_value value = {
		.value_data = label,
		.value_size = sizeof *label,
	};

	if (b_hashmap_get(c->s_labels, &key)) {
		return IVY_ERR_NAME_EXISTS;
	}

	b_status status = b_hashmap_put(c->s_labels, &key, &value);

	return ivy_status_from_b_status(status);
}

static enum ivy_status put_label_ref(
	struct ivy_assembler *as, struct assembler_scope *scope,
	const struct ivy_asm_token *label_name, size_t ref_offset)
{
	struct block_assembler_scope *c = (struct block_assembler_scope *)scope;
	struct label *label = label_create(LABEL_UNRESOLVED);

	label->l_name = label_name;
	label->l_offset = ref_offset;

	b_queue_push_back(&c->s_label_refs, &label->l_entry);
	return IVY_OK;
}

const struct assembler_scope_type block_scope_type = {
	.s_scope_size = sizeof(struct block_assembler_scope),
	.s_init_scope = init_scope,
	.s_finish_scope = finish_scope,
	.s_put_instr = put_instr,
	.s_put_label = put_label,
	.s_put_label_ref = put_label_ref,
};