#include "assembler.h"

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

struct constpool_assembler_scope {
	struct assembler_scope s_base;
	b_dict *s_strings;
	size_t s_next_slot;
};

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

	assembler_write_data(as, &header, sizeof header);

	struct constpool_assembler_scope *c
		= (struct constpool_assembler_scope *)scope;
	c->s_strings = b_dict_create();
	c->s_next_slot = 0;

	return IVY_OK;
}

static ivy_extended_data_key get_cached_string(
	struct constpool_assembler_scope *scope, const char *s)
{
	b_number *key = b_dict_at(scope->s_strings, s);
	if (!key) {
		return IVY_EX_DATA_KEY_NULL;
	}

	return (ivy_extended_data_key)b_number_get_int32(key);
}

static void put_cached_string(
	struct constpool_assembler_scope *scope, const char *s,
	ivy_extended_data_key key)
{
	b_dict_put(scope->s_strings, s, B_RV_INT32(key));
}

static ivy_extended_data_key write_string(struct ivy_assembler *as, const char *s)
{
	struct constpool_assembler_scope *scope
		= (struct constpool_assembler_scope *)assembler_get_scope(as);

	ivy_extended_data_key key = get_cached_string(scope, s);
	if (key != IVY_EX_DATA_KEY_NULL) {
		return key;
	}

	size_t len = strlen(s);

	struct ivy_bin_string str = {0};
	str.s_hash = b_i32_htob((uint32_t)b_hash_cstr(s));
	str.s_len = b_i32_htob((uint32_t)len);

	key = assembler_write_extended_data(as, &str, sizeof str);

	assembler_write_extended_data(as, s, len);
	put_cached_string(scope, s, key);

	return key;
}

static ivy_extended_data_key write_selector(
	struct ivy_assembler *as, const struct ivy_selector *sel)
{
	struct ivy_bin_selector dat = {0};
	switch (sel->sel_recipient) {
	case IVY_SEL_CLASS:
		dat.sel_flags = IVY_BIN_SELECTOR_CLASS;
		break;
	case IVY_SEL_OBJECT:
		dat.sel_flags = IVY_BIN_SELECTOR_OBJECT;
		break;
	default:
		return IVY_EX_DATA_KEY_NULL;
	}

	if (sel->sel_name) {
		dat.sel_name = b_i32_htob(write_string(as, sel->sel_name));
	}

	size_t i = 0;
	size_t nr_args = b_queue_length(&sel->sel_args);
	dat.sel_nr_args = (uint8_t)nr_args;
	/* TODO hash. */

	ivy_extended_data_key *arg_handles
		= calloc(nr_args, sizeof(ivy_extended_data_key));

	b_queue_entry *entry = b_queue_first(&sel->sel_args);
	while (entry) {
		struct ivy_selector_arg *arg
			= b_unbox(struct ivy_selector_arg, entry, arg_entry);
		arg_handles[i++] = write_string(as, arg->arg_label);
		entry = b_queue_next(entry);
	}

	ivy_extended_data_key selector_handle
		= assembler_write_extended_data(as, &dat, sizeof dat);
	for (i = 0; i < nr_args; i++) {
		b_i32 arg_handle = b_i32_htob(arg_handles[i]);
		assembler_write_extended_data(as, &arg_handle, sizeof arg_handle);
	}

	free(arg_handles);
	return selector_handle;
}

static ivy_extended_data_key write_ident(
	struct ivy_assembler *as, const struct ivy_ident *id)
{
	struct ivy_bin_ident dat = {0};
	/* TODO hash. */

	size_t nr_parts = b_queue_length(&id->id_parts);
	size_t i = 0;

	ivy_extended_data_key *part_handles
		= calloc(nr_parts, sizeof(ivy_extended_data_key));
	b_queue_entry *entry = b_queue_first(&id->id_parts);
	while (entry) {
		struct ivy_ident_part *arg
			= b_unbox(struct ivy_ident_part, entry, p_entry);
		part_handles[i++] = write_string(as, arg->p_str);
		entry = b_queue_next(entry);
	}

	dat.id_nr_parts = (uint8_t)nr_parts;

	ivy_extended_data_key ident_handle
		= assembler_write_extended_data(as, &dat, sizeof dat);

	for (i = 0; i < nr_parts; i++) {
		b_i32 part_handle = b_i32_htob(part_handles[i]);
		assembler_write_extended_data(as, &part_handle, sizeof part_handle);
	}

	free(part_handles);
	return ident_handle;
}

static enum ivy_status put_pval(
	struct ivy_assembler *as, struct assembler_scope *scope,
	enum ivy_assembler_pval_type type, size_t slot, const void *val,
	ivy_extended_data_key *key)
{
	struct ivy_bin_constpool_table_entry entry = {0};
	uintptr_t i = *(uintptr_t *)val;
	ivy_extended_data_key k = IVY_EX_DATA_KEY_NULL;
	struct constpool_assembler_scope *s
		= (struct constpool_assembler_scope *)assembler_get_scope(as);

	if (slot < s->s_next_slot) {
		return IVY_ERR_NAME_EXISTS;
	}

	switch (type) {
	case IVY_ASM_PVAL_STRING:
		entry.e_type = b_i32_htob(IVY_CONSTPOOL_TABLE_STRING);
		k = write_string(as, val);
		entry.e_ex_handle = b_i32_htob(k);
		break;
	case IVY_ASM_PVAL_IDENT:
		entry.e_type = b_i32_htob(IVY_CONSTPOOL_TABLE_IDENT);
		k = write_ident(as, val);
		entry.e_ex_handle = b_i32_htob(k);
		break;
	case IVY_ASM_PVAL_ATOM:
		entry.e_type = b_i32_htob(IVY_CONSTPOOL_TABLE_ATOM);
		k = write_string(as, val);
		entry.e_ex_handle = b_i32_htob(k);
		break;
	case IVY_ASM_PVAL_SINT:
		entry.e_type = b_i32_htob(IVY_CONSTPOOL_TABLE_INT);
		entry.e_int = b_i32_htob((uint32_t)i);
		break;
	case IVY_ASM_PVAL_UINT:
		entry.e_type = b_i32_htob(IVY_CONSTPOOL_TABLE_UINT);
		entry.e_int = b_i32_htob((uint32_t)i);
		break;
	case IVY_ASM_PVAL_SELECTOR:
		entry.e_type = b_i32_htob(IVY_CONSTPOOL_TABLE_SELECTOR);
		k = write_selector(as, val);
		entry.e_ex_handle = b_i32_htob(k);
		break;
	default:
		return IVY_ERR_NOT_SUPPORTED;
	}

	if (key) {
		*key = k;
	}

	for (size_t i = s->s_next_slot; i < slot; i++) {
		struct ivy_bin_constpool_table_entry null_entry = {0};
		assembler_write_data(as, &null_entry, sizeof null_entry);
	}

	assembler_write_data(as, &entry, sizeof entry);
	s->s_next_slot = slot + 1;
	return IVY_OK;
}

const struct assembler_scope_type constpool_scope_type = {
	.s_scope_size = sizeof(struct constpool_assembler_scope),
	.s_init_scope = init_scope,
	.s_put_pval = put_pval,
};