#include <ivy/asm/assembler.h>
#include <ivy/asm/bin.h>
#include <blue/core/queue.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

#include "assembler.h"

extern struct assembler_scope_type class_scope_type;
extern struct assembler_scope_type constpool_scope_type;

static const struct assembler_scope_type *scope_types[] = {
	[IVY_ASM_SCOPE_CLASS] = &class_scope_type,
	[IVY_ASM_SCOPE_CONSTPOOL] = &constpool_scope_type,
};
static const size_t nr_scope_types = sizeof scope_types / sizeof scope_types[0];

struct asm_table_entry {
	b_queue_entry e_entry;
	struct ivy_bin_table_entry e_data;
};

struct ivy_assembler {
	struct assembler_scope *as_scope;
	b_queue as_table;
	
	FILE *as_data;
	size_t as_data_offset;

	FILE *as_xdat;
	extended_data_key as_next_xdata_key;
};

enum ivy_status ivy_assembler_create(FILE *fp, struct ivy_assembler **as)
{
	struct ivy_assembler *out = malloc(sizeof *out);
	
	if (!out) {
		return IVY_ERR_NO_MEMORY;
	}

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

	out->as_data = fp;
	out->as_xdat = tmpfile();

	struct ivy_bin_header header = {0};
	header.h_magic = b_i32_htob(0xAABBCCDD);
	
	fseek(fp, 0, SEEK_SET);
	fwrite(&header, sizeof header, 1, fp);
	out->as_data_offset = sizeof header;

	char blank[16] = {0};
	assembler_write_extended_data(out, blank, sizeof blank);

	*as = out;
	return IVY_OK;
}

static void pad(struct ivy_assembler *as, unsigned int boundary)
{
	while ((as->as_data_offset % boundary) != 0) {
		fputc(0, as->as_data);
		as->as_data_offset++;
	}
}

enum ivy_status ivy_assembler_finish(struct ivy_assembler *as)
{
	const size_t buf_len = 4096;
	size_t xdat_len = as->as_next_xdata_key;
	
	fseek(as->as_xdat, 0, SEEK_SET);
	unsigned char *buf = malloc(buf_len);

	if (!buf) {
		return IVY_ERR_NO_MEMORY;
	}

	pad(as, 16);

	struct ivy_bin_table_entry xdat = {0};
	xdat.e_offset = b_i64_htob(as->as_data_offset);
	xdat.e_size = b_i32_htob(xdat_len);
	xdat.e_type = b_i32_htob(IVY_TABLE_XDAT);

	while (1) {
		size_t r = fread(buf, 1, buf_len, as->as_xdat);
		fwrite(buf, 1, r, as->as_data);
		as->as_data_offset += r;

		if (r == buf_len) {
			continue;
		}

		if (ferror(as->as_xdat)) {
			free(buf);
			return IVY_ERR_IO_FAILURE;
		}

		break;
	}

	free(buf);

	unsigned short nr_table_entries = 0;

	struct ivy_bin_header header = {0};
	header.h_table_offset = b_i64_htob(as->as_data_offset);
	header.h_table_len = b_i16_htob(0);
	header.h_magic = b_i32_htob(IVY_BIN_MAGIC);

	pad(as, 16);

	b_queue_iterator it = {0};
	b_queue_foreach (&it, &as->as_table) {
		struct asm_table_entry *e = b_unbox(struct asm_table_entry, it.entry, e_entry);
		size_t w = fwrite(&e->e_data, 1, sizeof e->e_data, as->as_data);

		if (w < sizeof e->e_data) {
			return IVY_ERR_IO_FAILURE;
		}

		nr_table_entries++;
	}

	fwrite(&xdat, 1, sizeof xdat, as->as_data);
	nr_table_entries++;

	header.h_table_len = b_i16_htob(nr_table_entries);

	fseek(as->as_data, 0, SEEK_SET);
	fwrite(&header, 1, sizeof header, as->as_data);

	fclose(as->as_xdat);
	free(as);

	return IVY_OK;
}

size_t assembler_write_data(struct ivy_assembler *as, const void *p, size_t len)
{
	size_t offset = as->as_data_offset;

	fwrite(p, 1, len, as->as_data);

	as->as_data_offset += len;
	return offset;
}

extended_data_key assembler_write_extended_data(struct ivy_assembler *as, const void *p, size_t len)
{
	extended_data_key key = as->as_next_xdata_key;

	fwrite(p, 1, len, as->as_xdat);

	as->as_next_xdata_key += len;
	return key;
}

struct assembler_scope *assembler_get_scope(struct ivy_assembler *as)
{
	return as->as_scope;
}

enum ivy_status ivy_assembler_begin_scope(
	struct ivy_assembler *as, enum ivy_assembler_scope_type type,
	ivy_assembler_attrib_table attrib)
{
	if (type < 0 || type >= nr_scope_types) {
		return IVY_ERR_NOT_SUPPORTED;
	}

	const struct assembler_scope_type *type_info = scope_types[type];
	if (!type_info) {
		return IVY_ERR_NOT_SUPPORTED;
	}

	if (as->as_scope) {
		return IVY_ERR_NOT_SUPPORTED;
	}

	struct assembler_scope *scope = malloc(type_info->s_scope_size);
	if (!scope) {
		return IVY_ERR_NO_MEMORY;
	}

	memset(scope, 0x0, type_info->s_scope_size);

	while ((as->as_data_offset % 16) != 0) {
		fputc(0, as->as_data);
		as->as_data_offset++;
	}

	scope->s_ops = type_info;
	scope->s_type = type;
	scope->s_start_offset = as->as_data_offset;
	as->as_scope = scope;

	if (type_info->s_init_scope) {
		type_info->s_init_scope(as, scope, attrib);
	}

	return IVY_OK;
}

enum ivy_status ivy_assembler_end_scope(struct ivy_assembler *as)
{
	if (!as->as_scope) {
		return IVY_ERR_NOT_SUPPORTED;
	}

	const struct assembler_scope_type *type_info = as->as_scope->s_ops;
	if (!type_info) {
		return IVY_ERR_NOT_SUPPORTED;
	}

	if (type_info->s_finish_scope) {
		type_info->s_finish_scope(as, as->as_scope);
	}

	struct asm_table_entry *entry = malloc(sizeof *entry);
	if (!entry) {
		return IVY_ERR_NO_MEMORY;
	}

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

	entry->e_data.e_offset = b_i64_htob(as->as_scope->s_start_offset);
	entry->e_data.e_size = b_i32_htob(as->as_data_offset - as->as_scope->s_start_offset);

	switch (as->as_scope->s_type) {
	case IVY_ASM_SCOPE_CLASS:
		entry->e_data.e_type = b_i32_htob(IVY_TABLE_CLASS);
		break;
	case IVY_ASM_SCOPE_LAMBDA:
		entry->e_data.e_type = b_i32_htob(IVY_TABLE_LAMBDA);
		break;
	case IVY_ASM_SCOPE_MSGH:
		entry->e_data.e_type = b_i32_htob(IVY_TABLE_MSGH);
		break;
	case IVY_ASM_SCOPE_CONSTPOOL:
		entry->e_data.e_type = b_i32_htob(IVY_TABLE_POOL);
		break;
	case IVY_ASM_SCOPE_IMPORT:
		entry->e_data.e_type = b_i32_htob(IVY_TABLE_IMPORT);
		break;
	default:
		break;
	}

	b_queue_push_back(&as->as_table, &entry->e_entry);

	free(as->as_scope);
	as->as_scope = NULL;

	return IVY_OK;
}

enum ivy_status ivy_assembler_put_pval(
	struct ivy_assembler *as, enum ivy_assembler_pval_type type,
	unsigned long index, const void *val)
{
	struct assembler_scope *scope = as->as_scope;

	if (!scope) {
		return IVY_ERR_NOT_SUPPORTED;
	}

	const struct assembler_scope_type *type_info = scope->s_ops;
	if (!type_info || !type_info->s_put_pval) {
		return IVY_ERR_NOT_SUPPORTED;
	}

	return type_info->s_put_pval(as, scope, type, index, val);
}

enum ivy_status ivy_assembler_put_xval(
	struct ivy_assembler *as, enum ivy_assembler_xval_type type,
	const ivy_assembler_attrib_table attrib)
{
	struct assembler_scope *scope = as->as_scope;

	if (!scope) {
		return IVY_ERR_NOT_SUPPORTED;
	}

	const struct assembler_scope_type *type_info = scope->s_ops;
	if (!type_info || !type_info->s_put_xval) {
		return IVY_ERR_NOT_SUPPORTED;
	}

	return type_info->s_put_xval(as, scope, type, attrib);
}

enum ivy_status ivy_assembler_put_instr(
	struct ivy_assembler *as, const struct ivy_instr *i)
{
	struct assembler_scope *scope = as->as_scope;

	if (!scope) {
		return IVY_ERR_NOT_SUPPORTED;
	}

	const struct assembler_scope_type *type_info = scope->s_ops;
	if (!type_info || !type_info->s_put_instr) {
		return IVY_ERR_NOT_SUPPORTED;
	}

	return type_info->s_put_instr(as, scope, i);
}