#include "wasm_compile.h"
#include "stack_machine_instruction.h"

/* WebAssembly opcodes */
#define WASM_END           0x0B
#define WASM_CALL          0x10
#define WASM_LOCAL_GET     0x20
#define WASM_I32_LOAD      0x28
#define WASM_I32_LOAD8_S   0x2C
#define WASM_I32_LOAD8_U   0x2D
#define WASM_I32_LOAD16_S  0x2E
#define WASM_I32_LOAD16_U  0x2F
#define WASM_I32_STORE     0x36
#define WASM_I32_STORE8    0x3A
#define WASM_I32_STORE16   0x3B
#define WASM_I64_STORE32   0x3E
#define WASM_I32_CONST     0x41
#define WASM_I32_SUB       0x6B

int translate(FILE *handle, struct function *function, struct module *module)
{
	struct instruction *expr = NULL;
	uint32_t args_count = function->type->arguments_count;
	uint32_t locals_count = function->locals_count;
	uint32_t all_locals_count = args_count + locals_count;
	size_t i;
	int wasm_opcode;
	int matched;

	if (locals_count + (uint64_t) args_count > STACK_TOP_ADDR * 4) {
		PRERR("Too many locals in a function\n");
		goto fail;
	}

	for (i = locals_count + 3; i; i--) {
		if (i_const(im(0), &expr))
			goto fail;
	}

	/* function prologue */
	if (i_get_frame(                             &expr) ||
	    i_tee      (                             &expr) ||
	    i_load     (im(STACK_FRAME_BACKUP_ADDR), &expr) ||
	    i_store_p  (im(0x0),                     &expr) ||
	    i_store    (im(STACK_FRAME_BACKUP_ADDR), &expr))
		goto fail;

	/* actual function body */
	i = 0;
	while (1) {
		matched = 0;

		wasm_opcode = fgetc(handle);

		if (wasm_opcode == EOF) {
			PRERR(MSG_EOF);
			goto fail;
		}

		// TODO: make a function for each instruction type,
		// call them through some table...
		if (wasm_opcode == WASM_I32_SUB) {
			if (i_sub(&expr))
				goto fail;

			matched = 1;
		} else if (wasm_opcode == WASM_CALL) {
			uint32_t funcidx;
			struct instruction **target;

			if (leb_u32(handle, &funcidx)) {
				PRERR(MSG_BAD_NUM);
				goto fail;
			}

			if (funcidx >= module->functions_count) {
				PRERR(MSG_BAD_IDX("funcidx"));
				goto fail;
			}

			target = &module->functions[funcidx].translated_body;

			if (i_call(ptr(target), &expr))
				goto fail;

			matched = 1;
		} else if (wasm_opcode <= WASM_I64_STORE32 &&
			   wasm_opcode >= WASM_I32_LOAD) {
			uint32_t align, offset;

			if (leb_u32(handle, &align) ||
			    leb_u32(handle, &offset)) {
				PRERR(MSG_BAD_NUM);
				goto fail;
			}

			offset += MEMORY_BOTTOM_ADDR;

			// TODO: rewrite it some cleaner way
#define TRY(opcode, instr)					  \
			if (wasm_opcode == opcode) {		  \
				matched = 1;			  \
								  \
				if (i_##instr(im(offset), &expr)) \
					goto fail;		  \
			}

			TRY(WASM_I32_LOAD,     load_p);
			TRY(WASM_I32_LOAD8_S,  loadbsx_p);
			TRY(WASM_I32_LOAD8_U,  loadbzx_p);
			TRY(WASM_I32_LOAD16_S, loadwsx_p);
			TRY(WASM_I32_LOAD16_U, loadwzx_p);
			TRY(WASM_I32_STORE,    store_p);
			TRY(WASM_I32_STORE8,   storeb_p);
			TRY(WASM_I32_STORE16,  storew_p);
		} else if (wasm_opcode == WASM_LOCAL_GET) {
			uint32_t localidx;
			uint64_t offset_on_frame;

			if (leb_u32(handle, &localidx)) {
				PRERR(MSG_BAD_NUM);
				goto fail;
			}

			if (localidx >= all_locals_count) {
				PRERR(MSG_BAD_IDX("localidx"));
				goto fail;
			}

			offset_on_frame = all_locals_count - localidx + 1;

			if (localidx >= args_count)
				offset_on_frame -= 1;

			if (i_load  (im(STACK_FRAME_BACKUP_ADDR), &expr) ||
			    i_load_p(im(4 * offset_on_frame),     &expr))
				goto fail;

			matched = 1;
		} else if (wasm_opcode == WASM_I32_CONST) {
			uint32_t constant;

			if (leb_u32(handle, &constant)) {
				PRERR(MSG_BAD_NUM);
				goto fail;
			}

			if (i_const(im(constant), &expr))
				goto fail;

			matched = 1;
		} else if (wasm_opcode == WASM_END) {
			break;
		}

		if (!matched) {
			PRERR("Unknown Wasm opcode: %02x\n", wasm_opcode);
			goto fail;
		}
	}

        /* function epilogue */
	if (function->type->result) {
		if (i_load   (im(STACK_FRAME_BACKUP_ADDR),        &expr) ||
		    i_swap   (                                    &expr) ||
		    i_store_p(im(4 * (2 + all_locals_count - 1)), &expr) ||
		    i_load   (im(STACK_FRAME_BACKUP_ADDR),        &expr) ||
		    i_tee    (                                    &expr) ||
		    i_tee    (                                    &expr) ||
		    i_load_p (im(4 * (1 + locals_count)),         &expr) ||
		    i_store_p(im(4 * (2 + all_locals_count - 2)), &expr) ||
		    i_load_p (im(0),                              &expr) ||
		    i_store  (im(STACK_FRAME_BACKUP_ADDR),        &expr))
			goto fail;
	} else {
		/* It's a bit shorter if we don't return anything */
		if (i_load   (im(STACK_FRAME_BACKUP_ADDR),        &expr) ||
		    i_tee    (                                    &expr) ||
		    i_tee    (                                    &expr) ||
		    i_load_p (im(4 * (1 + locals_count)),         &expr) ||
		    i_store_p(im(4 * (2 + all_locals_count - 1)), &expr) ||
		    i_load_p (im(0),                              &expr) ||
		    i_store  (im(STACK_FRAME_BACKUP_ADDR),        &expr))
			goto fail;

	}

	i = locals_count + args_count + 2 + (function->type->result ? 0 : 1);

	while (i--) {
		if (i_drop(&expr))
			goto fail;
	}

	if (i_ret(&expr))
		goto fail;

	function->translated_body = expr;

	return 0;

fail:
	PRERR("Couldn't translate function to stack machine\n");

	free_expr(expr);

	return -1;
}