#include "reader.h"

#include <ctype.h>
#include <ivy/asm/bin.h>
#include <ivy/asm/reader.h>
#include <ivy/ident.h>
#include <ivy/selector.h>
#include <ivy/status.h>
#include <stdlib.h>
#include <string.h>

static enum ivy_status decode_header(
	const struct ivy_bin_header *hdr, struct ivy_asm_object_info *out)
{
	out->obj_magic = b_i32_btoh(hdr->h_magic);
	out->obj_nr_sections = b_i16_btoh(hdr->h_table_len);
	out->obj_table_offset = b_i64_btoh(hdr->h_table_offset);

	return IVY_OK;
}

static enum ivy_status decode_table_entry(
	const struct ivy_bin_table_entry *entry, struct ivy_asm_section_info *out)
{
	out->s_type = b_i32_btoh(entry->e_type);
	out->s_offset = b_i64_btoh(entry->e_offset);
	out->s_length = b_i32_btoh(entry->e_size);

	return IVY_OK;
}

enum ivy_status ivy_asm_reader_open(FILE *in, struct ivy_asm_reader **out)
{
	struct ivy_asm_reader *reader = malloc(sizeof *reader);
	if (!reader) {
		return IVY_ERR_NO_MEMORY;
	}

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

	reader->r_fp = in;

	fseek(in, 0, SEEK_SET);

	struct ivy_bin_header hdr;
	size_t r = fread(&hdr, sizeof hdr, 1, in);

	if (r != 1) {
		free(reader);
		return IVY_ERR_IO_FAILURE;
	}

	enum ivy_status status = decode_header(&hdr, &reader->r_info);
	if (status != IVY_OK) {
		free(reader);
		return status;
	}

	if (reader->r_info.obj_magic != IVY_BIN_MAGIC) {
		free(reader);
		return IVY_ERR_BAD_FORMAT;
	}

	fseek(in, reader->r_info.obj_table_offset, SEEK_SET);

	size_t nr_sections = reader->r_info.obj_nr_sections;
	struct ivy_bin_table_entry *table = calloc(nr_sections, sizeof *table);
	if (!table) {
		free(reader);
		return IVY_ERR_NO_MEMORY;
	}

	r = fread(table, sizeof *table, nr_sections, in);
	if (r != nr_sections) {
		free(table);
		free(reader);
		return IVY_ERR_IO_FAILURE;
	}

	struct ivy_asm_section_info *sections
		= calloc(nr_sections, sizeof *sections);
	for (size_t i = 0; i < nr_sections; i++) {
		status = decode_table_entry(&table[i], &sections[i]);

		if (status != IVY_OK) {
			free(table);
			free(sections);
			free(reader);

			return status;
		}
	}

	free(table);

	reader->r_sections = sections;

	*out = reader;
	return IVY_OK;
}

enum ivy_status ivy_asm_reader_close(struct ivy_asm_reader *reader)
{
	free(reader);
	return IVY_OK;
}

const struct ivy_asm_object_info *ivy_asm_reader_get_info(
	struct ivy_asm_reader *reader)
{
	return &reader->r_info;
}

const struct ivy_asm_section_info *ivy_asm_reader_get_sections(
	struct ivy_asm_reader *reader)
{
	return reader->r_sections;
}

static size_t find_section(struct ivy_asm_reader *reader, uint32_t type, size_t index)
{
	if (type == IVY_ASM_SECTION_ANY_TYPE) {
		return index;
	}

	size_t nr_sections = reader->r_info.obj_nr_sections;
	for (size_t i = 0; i < nr_sections; i++) {
		if (reader->r_sections[i].s_type != type) {
			continue;
		}

		if (index > 0) {
			index--;
			continue;
		}

		return i;
	}

	return (size_t)-1;
}

enum ivy_status ivy_asm_reader_open_section(
	struct ivy_asm_reader *reader, uint32_t type, size_t index,
	struct ivy_asm_section_reader **out)
{
	index = find_section(reader, type, index);
	if (index == (size_t)-1) {
		return IVY_ERR_NO_ENTRY;
	}

	if (index >= reader->r_info.obj_nr_sections) {
		return IVY_ERR_NO_ENTRY;
	}

	struct ivy_asm_section_reader *sect = malloc(sizeof *sect);
	if (!sect) {
		return IVY_ERR_NO_MEMORY;
	}

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

	sect->r_parent = reader;
	sect->r_info = &reader->r_sections[index];

	*out = sect;
	return IVY_OK;
}

enum ivy_status ivy_asm_reader_open_constpool(
	struct ivy_asm_reader *reader, size_t constpool_index,
	size_t xdat_index, struct ivy_asm_constpool_reader **out)
{
	struct ivy_asm_section_reader *pool = NULL, *xdat = NULL;
	enum ivy_status status = ivy_asm_reader_open_section(
		reader, IVY_TABLE_POOL, constpool_index, &pool);
	if (status == IVY_ERR_NO_ENTRY) {
		return status;
	}

	status = ivy_asm_reader_open_section(
		reader, IVY_TABLE_XDAT, xdat_index, &xdat);
	if (status == IVY_ERR_NO_ENTRY) {
		ivy_asm_section_reader_close(pool);
		return status;
	}

	struct ivy_asm_constpool_reader *pool_reader = malloc(sizeof *pool_reader);
	if (!pool_reader) {
		ivy_asm_section_reader_close(pool);
		ivy_asm_section_reader_close(xdat);
		return IVY_ERR_NO_MEMORY;
	}

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

	pool_reader->r_pool = pool;
	pool_reader->r_xdat = xdat;
	pool_reader->r_nr_entries
		= (pool->r_info->s_length - sizeof(struct ivy_bin_constpool))
	        / sizeof(struct ivy_bin_constpool_table_entry);

	*out = pool_reader;
	return IVY_OK;
}

const struct ivy_asm_section_info *ivy_asm_section_reader_get_info(
	struct ivy_asm_section_reader *reader)
{
	return reader->r_info;
}

enum ivy_status ivy_asm_section_reader_close(struct ivy_asm_section_reader *reader)
{
	free(reader);
	return IVY_OK;
}

enum ivy_status ivy_asm_section_reader_read(
	struct ivy_asm_section_reader *reader, size_t offset, size_t count,
	void *p, size_t *nr_read)
{
	size_t object_offset = reader->r_info->s_offset + offset;
	FILE *fp = reader->r_parent->r_fp;
	fseek(fp, object_offset, SEEK_SET);
	size_t r = fread(p, 1, count, fp);

	*nr_read = r;

	return ferror(fp) ? IVY_ERR_IO_FAILURE : IVY_OK;
}

size_t ivy_asm_constpool_reader_get_nr_entries(
	struct ivy_asm_constpool_reader *reader)
{
	return reader->r_nr_entries;
}

static enum ivy_status read_string_xdata(
	struct ivy_asm_constpool_reader *pool, size_t key, char **out)
{
	if (!key) {
		*out = NULL;
		return IVY_OK;
	}

	size_t r;
	struct ivy_bin_string str;

	enum ivy_status status = ivy_asm_section_reader_read(
		pool->r_xdat, key, sizeof str, &str, &r);

	if (status != IVY_OK) {
		return status;
	}

	if (r != sizeof str) {
		return IVY_ERR_BAD_FORMAT;
	}

	size_t str_len = b_i32_btoh(str.s_len);
	if (str_len == 0) {
		*out = NULL;
		return IVY_OK;
	}

	char *s = malloc(str_len + 1);
	if (!s) {
		return IVY_ERR_NO_MEMORY;
	}

	status = ivy_asm_section_reader_read(
		pool->r_xdat, key + sizeof str, str_len, s, &r);
	if (status != IVY_OK) {
		free(s);
		return status;
	}

	if (r != str_len) {
		free(s);
		return IVY_ERR_BAD_FORMAT;
	}

	s[str_len] = 0;
	*out = s;
	return IVY_OK;
}

enum ivy_status ivy_asm_constpool_reader_read_value(
	struct ivy_asm_constpool_reader *reader, size_t index,
	struct ivy_asm_constpool_value **out)
{
	struct ivy_bin_constpool_table_entry entry;
	enum ivy_status status = IVY_OK;
	size_t r, offset
		  = sizeof(struct ivy_bin_constpool) + (index * sizeof entry);

	status = ivy_asm_section_reader_read(
		reader->r_pool, offset, sizeof entry, &entry, &r);

	if (status != IVY_OK) {
		return status;
	}

	if (r != sizeof entry) {
		return IVY_ERR_BAD_FORMAT;
	}

	unsigned long type = b_i32_btoh(entry.e_type);

	struct ivy_asm_constpool_value *value = malloc(sizeof *value);
	if (!value) {
		return IVY_ERR_NO_MEMORY;
	}

	memset(value, 0x0, sizeof *value);
	size_t handle;

	switch (type) {
	case IVY_CONSTPOOL_TABLE_STRING: {
		value->v_type = IVY_ASM_CONSTPOOL_TYPE_STRING;
		handle = b_i32_btoh(entry.e_ex_handle);
		status = read_string_xdata(reader, handle, &value->v_str);
		break;
	}
	case IVY_CONSTPOOL_TABLE_INT:
		value->v_type = IVY_ASM_CONSTPOOL_TYPE_INT;
		value->v_int = b_i32_btoh(entry.e_int);
		break;
	case IVY_CONSTPOOL_TABLE_UINT:
		value->v_type = IVY_ASM_CONSTPOOL_TYPE_UINT;
		value->v_uint = b_i32_btoh(entry.e_int);
		break;
	case IVY_CONSTPOOL_TABLE_ATOM:
		value->v_type = IVY_ASM_CONSTPOOL_TYPE_ATOM;
		handle = b_i32_btoh(entry.e_ex_handle);
		status = read_string_xdata(reader, handle, &value->v_str);
		break;
	case IVY_CONSTPOOL_TABLE_SELECTOR: {
		value->v_type = IVY_ASM_CONSTPOOL_TYPE_SELECTOR;
		handle = b_i32_btoh(entry.e_ex_handle);

		struct ivy_bin_selector sel_entry;
		status = ivy_asm_section_reader_read(
			reader->r_xdat, handle, sizeof sel_entry, &sel_entry, &r);

		if (status != IVY_OK) {
			break;
		}

		if (r != sizeof sel_entry) {
			status = IVY_ERR_BAD_FORMAT;
			break;
		}

		struct ivy_selector *sel;
		status = ivy_selector_create(&sel);
		if (status != IVY_OK) {
			break;
		}

		if (sel_entry.sel_flags & IVY_BIN_SELECTOR_OBJECT) {
			ivy_selector_set_recipient(sel, IVY_SEL_OBJECT);
		} else if (sel_entry.sel_flags & IVY_BIN_SELECTOR_CLASS) {
			ivy_selector_set_recipient(sel, IVY_SEL_CLASS);
		}

		size_t name_key = b_i32_btoh(sel_entry.sel_name);
		status = read_string_xdata(reader, name_key, &sel->sel_name);
		if (status != IVY_OK) {
			ivy_selector_destroy(sel);
			break;
		}

		handle += sizeof sel_entry;
		size_t nr_args = sel_entry.sel_nr_args;
		for (size_t i = 0; i < nr_args; i++) {
			b_i32 arg;
			status = ivy_asm_section_reader_read(
				reader->r_xdat, handle, sizeof arg, &arg, &r);

			if (status != IVY_OK) {
				break;
			}

			if (r != sizeof arg) {
				status = IVY_ERR_BAD_FORMAT;
				break;
			}

			size_t arg_key = b_i32_btoh(arg);
			char *arg_name = NULL;
			status = read_string_xdata(reader, arg_key, &arg_name);
			if (status != IVY_OK) {
				break;
			}

			ivy_selector_add_arg(sel, arg_name, NULL);
			free(arg_name);

			handle += sizeof arg;
		}

		if (status != IVY_OK) {
			ivy_selector_destroy(sel);
			break;
		}

		value->v_sel = sel;
		break;
	}
	case IVY_CONSTPOOL_TABLE_IDENT: {
		value->v_type = IVY_ASM_CONSTPOOL_TYPE_IDENT;
		handle = b_i32_btoh(entry.e_ex_handle);

		struct ivy_bin_ident id_entry;
		status = ivy_asm_section_reader_read(
			reader->r_xdat, handle, sizeof id_entry, &id_entry, &r);

		if (status != IVY_OK) {
			break;
		}

		if (r != sizeof id_entry) {
			status = IVY_ERR_BAD_FORMAT;
			break;
		}

		struct ivy_ident *id;
		id = ivy_ident_create();
		if (!id) {
			status = IVY_ERR_NO_MEMORY;
			break;
		}

		handle += sizeof id_entry;
		size_t nr_parts = id_entry.id_nr_parts;
		for (size_t i = 0; i < nr_parts; i++) {
			b_i32 part;
			status = ivy_asm_section_reader_read(
				reader->r_xdat, handle, sizeof part, &part, &r);

			if (status != IVY_OK) {
				break;
			}

			if (r != sizeof part) {
				status = IVY_ERR_BAD_FORMAT;
				break;
			}

			size_t part_key = b_i32_btoh(part);
			char *part_name = NULL;
			status = read_string_xdata(reader, part_key, &part_name);
			if (status != IVY_OK) {
				break;
			}

			ivy_ident_add_part(id, part_name);
			free(part_name);

			handle += sizeof part;
		}

		if (status != IVY_OK) {
			ivy_ident_destroy(id);
			break;
		}

		value->v_ident = id;
		break;
	}
	default:
		value->v_type = IVY_ASM_CONSTPOOL_TYPE_NONE;
		break;
	}

	if (status != IVY_OK) {
		free(value);
		value = NULL;
	}

	*out = value;
	return status;
}

enum ivy_status ivy_asm_constpool_reader_close(
	struct ivy_asm_constpool_reader *reader)
{
	ivy_asm_section_reader_close(reader->r_pool);
	ivy_asm_section_reader_close(reader->r_xdat);
	free(reader);

	return IVY_OK;
}

enum ivy_status ivy_asm_constpool_value_destroy(struct ivy_asm_constpool_value *v)
{
	switch (v->v_type) {
	case IVY_ASM_CONSTPOOL_TYPE_IDENT:
		ivy_ident_destroy(v->v_ident);
		break;
	case IVY_ASM_CONSTPOOL_TYPE_SELECTOR:
		ivy_selector_destroy(v->v_sel);
		break;
	case IVY_ASM_CONSTPOOL_TYPE_STRING:
		free(v->v_str);
		break;
	default:
		break;
	}

	free(v);

	return IVY_OK;
}

bool ivy_asm_section_type_to_string(uint32_t in, char out[5])
{
	b_i32 v = b_i32_htob(in);

	for (size_t i = 0; i < sizeof in; i++) {
		char c = v.i_bytes[i];
		if (!isgraph(c)) {
			return false;
		}

		out[i] = c;
	}

	out[4] = 0;
	return true;
}