#include "cmd.h"

#include <blue/cmd.h>
#include <blue/object/string.h>
#include <blue/term.h>
#include <ctype.h>
#include <errno.h>
#include <inttypes.h>
#include <ivy/asm/assembler.h>
#include <ivy/asm/instr.h>
#include <ivy/asm/lex.h>
#include <ivy/asm/parse.h>
#include <ivy/asm/reader.h>
#include <ivy/file.h>
#include <ivy/ident.h>
#include <ivy/selector.h>
#include <stdio.h>
#include <stdlib.h>

enum dump_flags {
	DUMP_INSTRUCTIONS = 0x01u,
};

enum {
	ARG_BIN_FILE,
	OPT_HEADER,
	OPT_SECTIONS,
	OPT_DUMP,
	OPT_CLASSES,
	OPT_CONSTPOOL,
};

static void dump_instruction(b_i32 x)
{
	struct ivy_instr instr = {};
	enum ivy_status status = ivy_instr_decode(x, &instr);

	if (status != IVY_OK) {
		printf("<invalid>");
		return;
	}

#define PRINT_CASE(code, str)                                                  \
	case (IVY_INSTR_##code):                                               \
		fputs(str, stdout);                                            \
		break

	const struct ivy_instr_definition *op = instr.i_op;

	switch (op->i_id) {
		PRINT_CASE(LDR, "ldr");
		PRINT_CASE(STR, "str");
		PRINT_CASE(PUSH, "push");
		PRINT_CASE(POP, "pop");
		PRINT_CASE(MSG, "msg");
		PRINT_CASE(ADD, "add");
		PRINT_CASE(SUB, "sub");
		PRINT_CASE(MUL, "mul");
		PRINT_CASE(DIV, "div");
		PRINT_CASE(C_EQ, "c.eq");
		PRINT_CASE(C_NE, "c.ne");
		PRINT_CASE(C_LT, "c.lt");
		PRINT_CASE(C_LE, "c.le");
		PRINT_CASE(C_GT, "c.gt");
		PRINT_CASE(C_GE, "c.ge");
		PRINT_CASE(BR, "br");
		PRINT_CASE(BR_T, "br.t");
		PRINT_CASE(BR_F, "br.f");
		PRINT_CASE(OB_C, "ob.c");
		PRINT_CASE(OB_E, "ob.e");
		PRINT_CASE(LAM_C, "lam.c");
		PRINT_CASE(IT_G, "it.g");
		PRINT_CASE(IT_N, "it.n");
		PRINT_CASE(IT_V, "it.v");
		PRINT_CASE(RET, "ret");
		PRINT_CASE(RET_N, "ret.n");
	default:
		fputs("<unknown>", stdout);
	}
#undef PRINT_CASE

	switch (op->i_opcode) {
	case IVY_OP_LDR_SP_REG:
		printf(" x%ld, [sp, x%ld]", instr.i_arg[0], instr.i_arg[1]);
		break;
	case IVY_OP_LDR_SP_CONST:
		printf(" x%ld, [sp, #%ld]", instr.i_arg[0], instr.i_arg[1]);
		break;
	case IVY_OP_LDR_BP_REG:
		printf(" x%ld, [bp, x%ld]", instr.i_arg[0], instr.i_arg[1]);
		break;
	case IVY_OP_LDR_BP_CONST:
		printf(" x%ld, [bp, #%ld]", instr.i_arg[0], instr.i_arg[1]);
		break;
	case IVY_OP_LDR_SELF_REG:
		printf(" x%ld, [self, x%ld]", instr.i_arg[0], instr.i_arg[1]);
		break;
	case IVY_OP_LDR_SELF_CONST:
		printf(" x%ld, [self, #%ld]", instr.i_arg[0], instr.i_arg[1]);
		break;
	case IVY_OP_LDR_POOL_REG:
		printf(" x%ld, [pool, x%ld]", instr.i_arg[0], instr.i_arg[1]);
		break;
	case IVY_OP_LDR_POOL_CONST:
		printf(" x%ld, [pool, #%ld]", instr.i_arg[0], instr.i_arg[1]);
		break;
	case IVY_OP_LDR_CONST:
		printf(" x%ld, #%ld", instr.i_arg[0], instr.i_arg[1]);
		break;

	case IVY_OP_STR_SP_REG:
		printf(" x%ld, [sp, x%ld]", instr.i_arg[0], instr.i_arg[1]);
		break;
	case IVY_OP_STR_SP_CONST:
		printf(" x%ld, [sp, #%ld]", instr.i_arg[0], instr.i_arg[1]);
		break;
	case IVY_OP_STR_BP_REG:
		printf(" x%ld, [bp, x%ld]", instr.i_arg[0], instr.i_arg[1]);
		break;
	case IVY_OP_STR_BP_CONST:
		printf(" x%ld, [bp, #%ld]", instr.i_arg[0], instr.i_arg[1]);
		break;
	case IVY_OP_STR_SELF_REG:
		printf(" x%ld, [self, x%ld]", instr.i_arg[0], instr.i_arg[1]);
		break;
	case IVY_OP_STR_SELF_CONST:
		printf(" x%ld, [self, #%ld]", instr.i_arg[0], instr.i_arg[1]);
		break;

	case IVY_OP_PUSH_REG:
	case IVY_OP_POP:
		printf(" x%ld", instr.i_arg[0]);
		break;
	case IVY_OP_PUSH_CONST:
		printf(" #%ld", instr.i_arg[0]);
		break;

	case IVY_OP_MSG_REG:
		printf(" x%ld, x%ld, x%lx", instr.i_arg[0], instr.i_arg[1],
		       instr.i_arg[2]);
		break;
	case IVY_OP_MSG_CONST:
		printf(" x%ld, x%ld, #%lx", instr.i_arg[0], instr.i_arg[1],
		       instr.i_arg[2]);
		break;

	case IVY_OP_ADD:
	case IVY_OP_SUB:
	case IVY_OP_MUL:
	case IVY_OP_DIV:
		printf(" x%ld, x%ld, x%lx", instr.i_arg[0], instr.i_arg[1],
		       instr.i_arg[2]);
		break;

	case IVY_OP_C_EQ:
	case IVY_OP_C_NE:
	case IVY_OP_C_LT:
	case IVY_OP_C_LE:
	case IVY_OP_C_GT:
	case IVY_OP_C_GE:
		printf(" x%ld, x%ld, x%lx", instr.i_arg[0], instr.i_arg[1],
		       instr.i_arg[2]);
		break;

	case IVY_OP_BR:
	case IVY_OP_BR_T:
	case IVY_OP_BR_F:
		printf(" $%04lXh", instr.i_arg[0]);
		break;

	case IVY_OP_OB_C_REG:
		break;
	case IVY_OP_OB_C_CONST:
		break;
	case IVY_OP_OB_E:
		break;

	case IVY_OP_LAM_C_REG:
		break;
	case IVY_OP_LAM_C_CONST:
		break;

	case IVY_OP_IT_G:
		break;
	case IVY_OP_IT_N:
		break;
	case IVY_OP_IT_V:
		break;

	case IVY_OP_RET:
		break;
	case IVY_OP_RET_N:
		break;
	default:
		break;
	}
}

static enum ivy_status dump_class(
	struct ivy_asm_reader *object, struct ivy_asm_constpool_reader *pool,
	size_t index)
{
	struct ivy_asm_section_reader *xcls = NULL;
	struct ivy_asm_constpool_value *class_ident = NULL;
	enum ivy_status status = ivy_asm_reader_open_section(
		object, IVY_TABLE_CLASS, index, &xcls);

	if (status != IVY_OK) {
		goto cleanup;
	}

	const struct ivy_asm_section_info *class_info
		= ivy_asm_section_reader_get_info(xcls);

	size_t r;
	struct ivy_bin_class class_header;
	status = ivy_asm_section_reader_read(
		xcls, 0, sizeof class_header, &class_header, &r);
	if (status != IVY_OK) {
		goto cleanup;
	}

	if (r != sizeof class_header) {
		status = IVY_ERR_BAD_FORMAT;
		goto cleanup;
	}

	status = ivy_asm_constpool_reader_read_value(
		pool, b_i32_btoh(class_header.c_ident), &class_ident);
	if (status != IVY_OK) {
		goto cleanup;
	}

	if (class_ident->v_type != IVY_ASM_CONSTPOOL_TYPE_IDENT) {
		status = IVY_ERR_BAD_FORMAT;
		goto cleanup;
	}

	char class_ident_str[128];
	ivy_ident_to_string(
		class_ident->v_ident, class_ident_str, sizeof class_ident_str);

	printf("  * %s\n", class_ident_str);

	struct ivy_bin_class_table_entry entry;
	size_t nr_entries
		= (class_info->s_length - sizeof class_header) / sizeof entry;
	size_t offset = sizeof class_header;

	for (size_t i = 0; i < nr_entries; i++) {
		status = ivy_asm_section_reader_read(
			xcls, offset, sizeof entry, &entry, &r);
		if (status != IVY_OK) {
			break;
		}

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

		printf("    [%03zu] ", i);

		char s[256];
		uint32_t type = b_i32_btoh(entry.e_type);
		switch (type) {
		case IVY_CLASS_TABLE_PROP: {
			printf("type:property\n");
			struct ivy_asm_constpool_value *prop_ident = NULL;
			status = ivy_asm_constpool_reader_read_value(
				pool, b_i32_btoh(entry.e_property.p_ident),
				&prop_ident);
			if (status != IVY_OK) {
				break;
			}

			if (prop_ident->v_type != IVY_ASM_CONSTPOOL_TYPE_IDENT) {
				status = IVY_ERR_BAD_FORMAT;
				break;
			}

			uint32_t get = b_i32_btoh(entry.e_property.p_get);
			uint32_t set = b_i32_btoh(entry.e_property.p_set);

			ivy_ident_to_string(prop_ident->v_ident, s, sizeof s);
			ivy_asm_constpool_value_destroy(prop_ident);
			printf("          ident:%s, get:", s);
			if (get) {
				printf("0x%" PRIx32 "[%" PRIu32 "]", get, get);
			} else {
				printf("<none>");
			}

			printf(", set:");

			if (set) {
				printf("0x%" PRIx32 "[%" PRIu32 "]", set, set);
			} else {
				printf("<none>");
			}

			break;
		}
		case IVY_CLASS_TABLE_MVAR: {
			printf("type:variable\n");
			struct ivy_asm_constpool_value *var_ident = NULL;
			status = ivy_asm_constpool_reader_read_value(
				pool, b_i32_btoh(entry.e_mvar.m_ident), &var_ident);
			if (status != IVY_OK) {
				break;
			}

			if (var_ident->v_type != IVY_ASM_CONSTPOOL_TYPE_IDENT) {
				status = IVY_ERR_BAD_FORMAT;
				break;
			}

			uint32_t index = b_i32_btoh(entry.e_mvar.m_index);
			ivy_ident_to_string(var_ident->v_ident, s, sizeof s);
			ivy_asm_constpool_value_destroy(var_ident);
			printf("          ident:%s, index:0x%" PRIx32
			       "[%" PRIu32 "]",
			       s, index, index);

			break;
		}
		case IVY_CLASS_TABLE_MSGH: {
			printf("type:method\n");
			struct ivy_asm_constpool_value *sel = NULL;
			status = ivy_asm_constpool_reader_read_value(
				pool, b_i32_btoh(entry.e_msgh.msg_selector), &sel);
			if (status != IVY_OK) {
				break;
			}

			if (sel->v_type != IVY_ASM_CONSTPOOL_TYPE_SELECTOR) {
				status = IVY_ERR_BAD_FORMAT;
				break;
			}

			uint32_t block = b_i32_btoh(entry.e_msgh.msg_block);
			ivy_selector_to_string(sel->v_sel, s, sizeof s);
			ivy_asm_constpool_value_destroy(sel);

			printf("          selector:[%s], block:0x%" PRIx32
			       "[%" PRIu32 "]",
			       s, block, block);

			break;
		}
		default:
			break;
		}

		printf("\n");
		offset += sizeof entry;
	}

cleanup:
	if (class_ident) {
		ivy_asm_constpool_value_destroy(class_ident);
	}

	if (xcls) {
		ivy_asm_section_reader_close(xcls);
	}

	return status;
}

static enum ivy_status dump_classes(struct ivy_asm_reader *object)
{
	struct ivy_asm_constpool_reader *pool;
	enum ivy_status status
		= ivy_asm_reader_open_constpool(object, 0, 0, &pool);
	if (status != IVY_OK) {
		return status;
	}

	printf("classes:\n");
	for (size_t i = 0;; i++) {
		status = dump_class(object, pool, i);
		if (status == IVY_ERR_NO_ENTRY) {
			break;
		}

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

	return IVY_OK;
}

static enum ivy_status dump_constpool(struct ivy_asm_reader *object)
{
	printf("constpool:\n");
	struct ivy_asm_constpool_reader *pool = NULL;
	enum ivy_status status
		= ivy_asm_reader_open_constpool(object, 0, 0, &pool);
	if (status == IVY_ERR_NO_ENTRY) {
		b_err("object file has no constpool.\n");
		return status;
	}

	size_t nr_entries = ivy_asm_constpool_reader_get_nr_entries(pool);

	for (size_t i = 0; i < nr_entries; i++) {
		struct ivy_asm_constpool_value *v;
		status = ivy_asm_constpool_reader_read_value(pool, i, &v);
		if (status != IVY_OK) {
			b_err("cannot read constpool value");
			b_i("reason: %s", ivy_status_to_string(status));
			break;
		}

		printf("  [%03zu] ", i);

		switch (v->v_type) {
		case IVY_ASM_CONSTPOOL_TYPE_STRING:
			printf("string \"%s\"", v->v_str);
			break;
		case IVY_ASM_CONSTPOOL_TYPE_SELECTOR: {
			char s[256];
			ivy_selector_to_string(v->v_sel, s, sizeof s);
			printf("selector %s", s);
			break;
		}
		case IVY_ASM_CONSTPOOL_TYPE_IDENT: {
			char s[256];
			ivy_ident_to_string(v->v_ident, s, sizeof s);
			printf("ident %s", s);
			break;
		}
		case IVY_ASM_CONSTPOOL_TYPE_INT:
			printf("sint #0x%08llx[%lld]", v->v_int, v->v_int);
			break;
		case IVY_ASM_CONSTPOOL_TYPE_UINT:
			printf("uint #0x%08llx[%llu]", v->v_uint, v->v_uint);
			break;
		case IVY_ASM_CONSTPOOL_TYPE_NONE:
			printf("<null>");
			break;
		case IVY_ASM_CONSTPOOL_TYPE_ATOM:
			printf("atom $%s", v->v_str);
			break;
		default:
			printf("<unknown>");
			break;
		}

		fputc('\n', stdout);
	}

	return status;
}

static void dump_instructions(b_i32 *x, size_t buffer_size)
{
	size_t nr_instr = buffer_size / sizeof *x;

	for (size_t i = 0; i < nr_instr; i++) {
		if (i > 0) {
			fputs("; ", stdout);
		}

		dump_instruction(x[i]);
	}
}

static void dump_strings(uint32_t *x, size_t buffer_size)
{
	char *c = (char *)x;

	for (int ii = 0; ii < buffer_size; ii++) {
		if (isgraph(c[ii])) {
			fputc(c[ii], stdout);
		} else {
			fputc('.', stdout);
		}
	}
}

static enum ivy_status dump_header(
	struct ivy_asm_section_reader *sect, size_t *dump_offset)
{
	const struct ivy_asm_section_info *info
		= ivy_asm_section_reader_get_info(sect);
	enum ivy_status status = IVY_OK;
	size_t r = 0;

	switch (info->s_type) {
	case IVY_TABLE_CLASS: {
		struct ivy_bin_class xcls;
		status = ivy_asm_section_reader_read(
			sect, 0, sizeof xcls, &xcls, &r);
		if (status != IVY_OK) {
			break;
		}

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

		printf("                ident=0x%04x\n", b_i32_btoh(xcls.c_ident));
		*dump_offset = sizeof xcls;
		break;
	}
	case IVY_TABLE_BLOCK: {
		struct ivy_bin_block text;
		status = ivy_asm_section_reader_read(
			sect, 0, sizeof text, &text, &r);
		if (status != IVY_OK) {
			break;
		}

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

		long index = b_i32_btoh(text.b_index);
		printf("                index=0x%04lx [%ld]\n", index, index);
		*dump_offset = sizeof text;
		break;
	}
	case IVY_TABLE_XDAT:
	default:
		*dump_offset = 0;
		break;
	}

	return status;
}

static enum ivy_status dump_section(
	struct ivy_asm_reader *reader, size_t index, size_t chunks_per_line,
	enum dump_flags flags)
{
	struct ivy_asm_section_reader *sect = NULL;
	enum ivy_status status = ivy_asm_reader_open_section(
		reader, IVY_ASM_SECTION_ANY_TYPE, index, &sect);
	if (status != IVY_OK) {
		return status;
	}

	const struct ivy_asm_section_info *info
		= ivy_asm_section_reader_get_info(sect);
	size_t len = info->s_length;
	size_t buffer_size = chunks_per_line * sizeof(uint32_t);
	uint32_t *x = malloc(buffer_size);
	size_t bytes_per_line = chunks_per_line * sizeof *x;

	size_t offset = 0;
	status = dump_header(sect, &offset);
	if (status != IVY_OK) {
		return status;
	}

	for (size_t i = offset; i < len; i += buffer_size) {
		size_t r;
		status = ivy_asm_section_reader_read(sect, i, buffer_size, x, &r);

		if (status != IVY_OK) {
			break;
		}

		if (r != buffer_size) {
			break;
		}

		printf("    %05zx  |  ", i);

		unsigned char *z = (unsigned char *)x;
		for (size_t ii = 0; ii < bytes_per_line; ii++) {
			if (ii > 0 && (ii % sizeof *x) == 0) {
				fputc(' ', stdout);
			}

			printf("%02x", z[ii]);
		}

		printf("  |  ");

		if (flags & DUMP_INSTRUCTIONS) {
			dump_instructions((b_i32 *)x, buffer_size);
		} else {
			dump_strings(x, buffer_size);
		}

		fputc('\n', stdout);
	}

	free(x);
	ivy_asm_section_reader_close(sect);
	return status;
}

static enum ivy_status dump_section_table(struct ivy_asm_reader *reader, bool dump)
{
	const struct ivy_asm_section_info *section_table
		= ivy_asm_reader_get_sections(reader);
	const struct ivy_asm_object_info *object_info
		= ivy_asm_reader_get_info(reader);

	printf("sections:\n");
	for (size_t i = 0; i < object_info->obj_nr_sections; i++) {
		const struct ivy_asm_section_info *sect = &section_table[i];
		char type[5];

		printf(" * #%03zu: ", i);

		if (ivy_asm_section_type_to_string(sect->s_type, type)) {
			printf("[%4s]", type);
		} else {
			printf("[%08" PRIx32 "]", sect->s_type);
		}

		printf(" offset=0x%04zx, length=0x%04zx\n", sect->s_offset,
		       sect->s_length);

		size_t chunks_per_line = 1;
		enum dump_flags flags = 0;

		switch (sect->s_type) {
		case IVY_TABLE_XDAT:
		case IVY_TABLE_CLASS:
			chunks_per_line = 4;
			break;
		case IVY_TABLE_POOL:
			chunks_per_line = 2;
			break;
		case IVY_TABLE_BLOCK:
			chunks_per_line = 1;
			flags |= DUMP_INSTRUCTIONS;
			break;
		default:
			chunks_per_line = 1;
			break;
		}

		if (dump) {
			dump_section(reader, i, chunks_per_line, flags);
			fputc('\n', stdout);
		}
	}
	return IVY_OK;
}

static int disassemble(
	const b_command *cmd, const b_arglist *args, const b_array *_)
{
	bool header = b_arglist_get_count(args, OPT_HEADER, B_COMMAND_INVALID_ID)
	            > 0;
	bool sections
		= b_arglist_get_count(args, OPT_SECTIONS, B_COMMAND_INVALID_ID) > 0;
	bool dump = b_arglist_get_count(args, OPT_DUMP, B_COMMAND_INVALID_ID) > 0;
	bool constpool
		= b_arglist_get_count(args, OPT_CONSTPOOL, B_COMMAND_INVALID_ID)
	        > 0;
	bool classes
		= b_arglist_get_count(args, OPT_CLASSES, B_COMMAND_INVALID_ID) > 0;

	const char *in_path = NULL;

	b_arglist_get_string(args, B_COMMAND_INVALID_ID, ARG_BIN_FILE, 0, &in_path);
	if (!in_path) {
		b_err("no input file specified.");
		return -1;
	}

	FILE *in = fopen(in_path, "rb");
	if (!in) {
		b_err("cannot open object file");
		b_i("reason: %s", strerror(errno));

		return -1;
	}

	struct ivy_asm_reader *reader = NULL;
	enum ivy_status status = ivy_asm_reader_open(in, &reader);

	if (status != IVY_OK) {
		b_err("cannot open object file");
		b_i("reason: %s", ivy_status_to_string(status));

		fclose(in);
		return -1;
	}

	const struct ivy_asm_object_info *object_info
		= ivy_asm_reader_get_info(reader);

	if (header) {
		printf("header:\n");
		printf("\tmagic:        0x%08" PRIx32 "\n", object_info->obj_magic);
		printf("\ttable offset: 0x%zx\n", object_info->obj_table_offset);
		printf("\tnr sections:  %zu\n", object_info->obj_nr_sections);
		printf("\n");
	}

	if (sections) {
		status = dump_section_table(reader, dump);
	}

	if (constpool) {
		status = dump_constpool(reader);
	}

	if (classes) {
		status = dump_classes(reader);
	}

	ivy_asm_reader_close(reader);

	return 0;
}

B_COMMAND(CMD_DISASSEMBLE, CMD_ROOT)
{
	B_COMMAND_NAME("disassemble");
	B_COMMAND_SHORT_NAME('D');
	B_COMMAND_DESC("disassemble an Ivy object file.");
	B_COMMAND_FLAGS(B_COMMAND_SHOW_HELP_BY_DEFAULT);
	B_COMMAND_FUNCTION(disassemble);

	B_COMMAND_ARG(ARG_BIN_FILE)
	{
		B_ARG_NAME("input file");
		B_ARG_DESC("the Ivy object file to disassemble.");
		B_ARG_NR_VALUES(1);
	}

	B_COMMAND_OPTION(OPT_HEADER)
	{
		B_OPTION_SHORT_NAME('h');
		B_OPTION_LONG_NAME("header");
		B_OPTION_DESC("print the object file header.");
	}

	B_COMMAND_OPTION(OPT_SECTIONS)
	{
		B_OPTION_SHORT_NAME('s');
		B_OPTION_LONG_NAME("section-table");
		B_OPTION_DESC("print the object file section table.");
	}

	B_COMMAND_OPTION(OPT_DUMP)
	{
		B_OPTION_SHORT_NAME('d');
		B_OPTION_LONG_NAME("dump");
		B_OPTION_DESC(
			"decode and print the contents of each object "
			"section.");
	}

	B_COMMAND_OPTION(OPT_CLASSES)
	{
		B_OPTION_SHORT_NAME('c');
		B_OPTION_LONG_NAME("classes");
		B_OPTION_DESC(
			"print the classes contained in the object "
			"file.");
	}

	B_COMMAND_OPTION(OPT_CONSTPOOL)
	{
		B_OPTION_SHORT_NAME('p');
		B_OPTION_LONG_NAME("pool-data");
		B_OPTION_DESC(
			"print the constant pool data in the object "
			"file.");
	}

	B_COMMAND_HELP_OPTION();
}