#include "parse.h"

#include <blue/core/hash.h>
#include <blue/core/stringstream.h>
#include <ctype.h>
#include <ivy/asm/assembler.h>
#include <ivy/asm/bin.h>
#include <ivy/asm/instr.h>
#include <ivy/ident.h>
#include <ivy/selector.h>
#include <stdlib.h>

#define HASH_SELF         0x2d19e518d40792b7
#define HASH_POOL         0x8c22f10da88b1083
#define HASH_SP           0x08d93e07b5793c56
#define HASH_BP           0x08a64407b54decef

#define REG_INDEX_INVALID ((unsigned long long)-1)

enum instr_component {
	INSTR_NONE = 0,
	INSTR_OPCODE,
	INSTR_OPCODE_DOT,
	INSTR_OPERAND,
	INSTR_OPERAND_SEPARATOR,
	INSTR_OPERAND_INDEX_LEFT,
	INSTR_OPERAND_INDEX_BASE,
	INSTR_OPERAND_INDEX_SEPARATOR,
	INSTR_OPERAND_INDEX_OFFSET,
	INSTR_OPERAND_INDEX_RIGHT,
};

enum index_base {
	INDEX_NONE = 0,
	INDEX_SELF,
	INDEX_POOL,
	INDEX_SP,
	INDEX_BP,
};

enum arg_type {
	ARG_NONE = 0,
	ARG_CONST,
	ARG_LABEL,
	ARG_REG,
	ARG_INDEX_REG,
	ARG_INDEX_CONST,
};

struct label {
	b_queue_entry l_entry;
	struct ivy_asm_token *l_name;
	unsigned long long l_offset;
};

struct arg {
	enum arg_type arg_type;
	b_queue_entry arg_entry;

	union {
		struct ivy_asm_token *arg_const;
		struct ivy_asm_token *arg_label;
		struct {
			struct ivy_asm_token *reg_token;
			unsigned long long reg_index;
		} arg_reg;
		struct {
			enum index_base index_base;
			struct ivy_asm_token *index_base_token;

			struct ivy_asm_token *index_offset;
			unsigned long long index_offset_reg;
		} arg_index_reg;

		struct {
			enum index_base index_base;
			struct ivy_asm_token *index_base_token;

			struct ivy_asm_token *index_offset;
		} arg_index_const;
	};
};

struct block_parser_state {
	struct parser_state s_base;
	unsigned int s_prev_token;
	enum instr_component s_prev_component;

	b_queue s_labels;

	b_queue s_mnemonic;

	b_queue s_args;
	struct arg *s_current_arg;
};

struct mnemonic {
	const char *m_name;
	uint64_t m_hash;
	enum ivy_instr_id m_id;
};

#define MNEMONIC(id, name, hash)                                               \
	{                                                                      \
		.m_id = (id), .m_name = (name), .m_hash = (hash)               \
	}

static const struct mnemonic mnemonics[] = {
	MNEMONIC(IVY_INSTR_LDR, "ldr", 0x127688191dd0471d),
	MNEMONIC(IVY_INSTR_STR, "str", 0x826e83195d0d60f0),
	MNEMONIC(IVY_INSTR_PUSH, "push", 0x6c80030e2762459d),
	MNEMONIC(IVY_INSTR_POP, "pop", 0x779b5819564f2f50),
	MNEMONIC(IVY_INSTR_MSG, "msg", 0x07e05a191745be82),
	MNEMONIC(IVY_INSTR_ADD, "add", 0xe70bc3190530e654),
	MNEMONIC(IVY_INSTR_SUB, "sub", 0x82719d195d0fc2f5),
	MNEMONIC(IVY_INSTR_MUL, "mul", 0x07e66519174a3ce1),
	MNEMONIC(IVY_INSTR_DIV, "div", 0xcaa83a18f46e5888),
	MNEMONIC(IVY_INSTR_C_EQ, "c.eq", 0x2d6e0e924e665476),
	MNEMONIC(IVY_INSTR_C_NE, "c.ne", 0x2d8d12924e810bd1),
	MNEMONIC(IVY_INSTR_C_LT, "c.lt", 0x2d861d924e7affd4),
	MNEMONIC(IVY_INSTR_C_LE, "c.le", 0x2d860e924e7ae657),
	MNEMONIC(IVY_INSTR_C_GT, "c.gt", 0x2d751f924e6c9007),
	MNEMONIC(IVY_INSTR_C_GE, "c.ge", 0x2d750e924e6c7324),
	MNEMONIC(IVY_INSTR_BR, "br", 0x08a64607b54df055),
	MNEMONIC(IVY_INSTR_BR_T, "br.t", 0x3e0bed9c06af59cf),
	MNEMONIC(IVY_INSTR_BR_F, "br.f", 0x3e0bdf9c06af4205),
	MNEMONIC(IVY_INSTR_OB_C, "ob.c", 0x70bdc0b4a72facf7),
	MNEMONIC(IVY_INSTR_OB_E, "ob.e", 0x70bdbeb4a72fa991),
	MNEMONIC(IVY_INSTR_LAM_C, "lam.c", 0xeb2c84ec62ed7472),
	MNEMONIC(IVY_INSTR_IT_G, "it.g", 0x2939c6c597a9d607),
	MNEMONIC(IVY_INSTR_IT_N, "it.n", 0x2939bdc597a9c6bc),
	MNEMONIC(IVY_INSTR_IT_V, "it.v", 0x2939b5c597a9b924),
	MNEMONIC(IVY_INSTR_RET, "ret", 0x89e9ae1960f4a6ec),
	MNEMONIC(IVY_INSTR_RET_N, "ret.n", 0x16b792f2c490f4d8),
};
static const size_t nr_mnemonics = sizeof mnemonics / sizeof mnemonics[0];

static unsigned long long get_register_index(struct ivy_asm_token *tok)
{
	if (tok->t_type != IVY_ASM_TOK_IDENT) {
		return REG_INDEX_INVALID;
	}

	const char *s = tok->t_str;
	if (tolower(s[0]) != 'x') {
		return REG_INDEX_INVALID;
	}

	char index_str[5] = {0};
	strncpy(index_str, s + 1, sizeof index_str - 1);
	char *ep = NULL;
	unsigned long long index = strtoul(index_str, &ep, 10);

	if (*ep || index >= 256) {
		return REG_INDEX_INVALID;
	}

	return index;
}

static enum index_base get_index_base(struct ivy_asm_token *tok)
{
	if (tok->t_type != IVY_ASM_TOK_IDENT) {
		return INDEX_NONE;
	}

	const char *s = tok->t_str;
	uint64_t hash = b_hash_string(s);

	switch (hash) {
	case HASH_SELF:
		if (!strcmp(s, "self")) {
			return INDEX_SELF;
		}

		return INDEX_NONE;
	case HASH_POOL:
		if (!strcmp(s, "pool")) {
			return INDEX_POOL;
		}

		return INDEX_NONE;
	case HASH_SP:
		if (!strcmp(s, "sp")) {
			return INDEX_SP;
		}

		return INDEX_NONE;
	case HASH_BP:
		if (!strcmp(s, "bp")) {
			return INDEX_BP;
		}

		return INDEX_NONE;
	default:
		return INDEX_NONE;
	}
}

static enum ivy_instr_id get_instruction_id(b_queue *mnemonic_tokens)
{
	char mnemonic[64];
	mnemonic[0] = 0;

	b_stringstream s;
	b_stringstream_begin(&s, mnemonic, sizeof mnemonic);

	unsigned int i = 0;

	b_queue_iterator it = {0};
	b_queue_iterator_begin(mnemonic_tokens, &it);
	while (b_queue_iterator_is_valid(&it)) {
		struct ivy_asm_token *tok
			= b_unbox(struct ivy_asm_token, it.entry, t_entry);
		b_queue_iterator_erase(&it);

		if (i > 0) {
			b_stringstream_add(&s, ".");
		}

		b_stringstream_add(&s, tok->t_str);
		i++;

		ivy_asm_token_destroy(tok);
	}

	uint64_t hash = b_hash_string(mnemonic);
	for (i = 0; i < nr_mnemonics; i++) {
		if (hash == mnemonics[i].m_hash
		    && !strcmp(mnemonic, mnemonics[i].m_name)) {
			return mnemonics[i].m_id;
		}
	}

	return IVY_INSTR_NONE;
}

static enum ivy_status write_instruction(
	struct ivy_asm_parser *p, struct block_parser_state *state)
{
	/* one more than the true maximum to catch if the input has specified
	 * too many operands */
#define MAX_ARGS 4

	unsigned int i = 0;

	enum ivy_instr_id id = get_instruction_id(&state->s_mnemonic);
	if (id == IVY_INSTR_NONE) {
		return IVY_ERR_BAD_SYNTAX;
	}

	enum ivy_instr_operand_type operand_types[MAX_ARGS] = {0};

	b_queue_iterator it = {0};
	b_queue_foreach (&it, &state->s_args) {
		struct arg *arg = b_unbox(struct arg, it.entry, arg_entry);

		if (i >= MAX_ARGS) {
			return IVY_ERR_BAD_SYNTAX;
		}

		switch (arg->arg_type) {
		case ARG_REG:
			operand_types[i++] = IVY_INSTR_OPERAND_REGISTER;
			break;
		case ARG_CONST:
		case ARG_LABEL:
			operand_types[i++] = IVY_INSTR_OPERAND_CONST;
			break;
		case ARG_INDEX_REG:
			switch (arg->arg_index_reg.index_base) {
			case INDEX_SELF:
				operand_types[i++]
					= IVY_INSTR_OPERAND_SELF_INDEX_REG;
				break;
			case INDEX_POOL:
				operand_types[i++]
					= IVY_INSTR_OPERAND_POOL_INDEX_REG;
			case INDEX_SP:
				operand_types[i++] = IVY_INSTR_OPERAND_SP_INDEX_REG;
			case INDEX_BP:
				operand_types[i++] = IVY_INSTR_OPERAND_BP_INDEX_REG;
			default:
				return IVY_ERR_BAD_SYNTAX;
			}
			break;
		case ARG_INDEX_CONST:
			switch (arg->arg_index_reg.index_base) {
			case INDEX_SELF:
				operand_types[i++]
					= IVY_INSTR_OPERAND_SELF_INDEX_CONST;
				break;
			case INDEX_POOL:
				operand_types[i++]
					= IVY_INSTR_OPERAND_POOL_INDEX_CONST;
				break;
			case INDEX_SP:
				operand_types[i++]
					= IVY_INSTR_OPERAND_SP_INDEX_CONST;
				break;
			case INDEX_BP:
				operand_types[i++]
					= IVY_INSTR_OPERAND_BP_INDEX_CONST;
				break;
			default:
				return IVY_ERR_BAD_SYNTAX;
			}
			break;
		default:
			return IVY_ERR_BAD_SYNTAX;
		}
	}

	const struct ivy_instr_definition *instr_info
		= ivy_instr_find(id, operand_types);
	if (!instr_info) {
		return IVY_ERR_BAD_SYNTAX;
	}

	struct ivy_instr instr = {0};
	instr.i_op = instr_info;
	i = 0;

	b_queue_iterator_begin(&state->s_args, &it);
	while (b_queue_iterator_is_valid(&it)) {
		struct arg *arg = b_unbox(struct arg, it.entry, arg_entry);
		b_queue_iterator_erase(&it);

		switch (arg->arg_type) {
		case ARG_REG:
			instr.i_arg[i++] = arg->arg_reg.reg_index;
			break;
		case ARG_CONST:
			instr.i_arg[i++] = arg->arg_const->t_int.sign
			                         ? arg->arg_const->t_int.v
			                         : arg->arg_const->t_int.uv;
			break;
		case ARG_LABEL:
			instr.i_arg[i++] = 0;
			ivy_assembler_put_label_ref(
				p->p_assembler, arg->arg_label,
				ivy_assembler_get_ptr(p->p_assembler));
			break;
		case ARG_INDEX_REG:
			instr.i_arg[i++] = arg->arg_index_reg.index_offset_reg;
			break;
		case ARG_INDEX_CONST:
			instr.i_arg[i++]
				= arg->arg_index_const.index_offset->t_int.sign
			                ? arg->arg_index_const.index_offset
			                          ->t_int.v
			                : arg->arg_index_const.index_offset
			                          ->t_int.uv;
			break;
		default:
			return IVY_ERR_BAD_SYNTAX;
		}

		free(arg);
	}

	ivy_assembler_put_instr(p->p_assembler, &instr);
	return IVY_OK;
}

static enum ivy_status push_const_arg(
	struct block_parser_state *state, struct ivy_asm_token *tok)
{
	struct arg *arg = malloc(sizeof *arg);
	if (!arg) {
		return IVY_ERR_NO_MEMORY;
	}

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

	arg->arg_type = ARG_CONST;
	arg->arg_const = tok;

	b_queue_push_back(&state->s_args, &arg->arg_entry);
	return IVY_OK;
}

static enum ivy_status push_label(
	struct ivy_asm_parser *p, struct block_parser_state *state,
	struct ivy_asm_token *tok)
{
	return ivy_assembler_put_label(
		p->p_assembler, tok, ivy_assembler_get_ptr(p->p_assembler));
}

static enum ivy_status push_label_arg(
	struct block_parser_state *state, struct ivy_asm_token *tok,
	unsigned long long reg_index)
{
	struct arg *arg = malloc(sizeof *arg);
	if (!arg) {
		return IVY_ERR_NO_MEMORY;
	}

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

	arg->arg_type = ARG_LABEL;
	arg->arg_label = tok;

	b_queue_push_back(&state->s_args, &arg->arg_entry);
	return IVY_OK;
}

static enum ivy_status push_reg_arg(
	struct block_parser_state *state, struct ivy_asm_token *tok,
	unsigned long long reg_index)
{
	struct arg *arg = malloc(sizeof *arg);
	if (!arg) {
		return IVY_ERR_NO_MEMORY;
	}

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

	arg->arg_type = ARG_REG;
	arg->arg_reg.reg_token = tok;
	arg->arg_reg.reg_index = reg_index;

	b_queue_push_back(&state->s_args, &arg->arg_entry);
	return IVY_OK;
}

static enum ivy_status parse_linefeed(
	struct ivy_asm_parser *ctx, struct ivy_asm_token *tok)
{
	struct block_parser_state *state
		= (struct block_parser_state *)asm_parser_get_state(ctx);

	switch (state->s_prev_component) {
	case INSTR_NONE:
		return IVY_OK;
	case INSTR_OPCODE:
	case INSTR_OPERAND:
	case INSTR_OPERAND_INDEX_RIGHT:
		state->s_prev_component = INSTR_NONE;
		return write_instruction(ctx, state);
	default:
		return IVY_ERR_BAD_SYNTAX;
	}
}

static enum ivy_status parse_int(
	struct ivy_asm_parser *ctx, struct ivy_asm_token *tok)
{
	struct block_parser_state *state
		= (struct block_parser_state *)asm_parser_get_state(ctx);

	switch (state->s_prev_component) {
	case INSTR_OPCODE:
	case INSTR_OPERAND_SEPARATOR:
		push_const_arg(state, tok);
		state->s_prev_component = INSTR_OPERAND;
		return IVY_OK;
	case INSTR_OPERAND_INDEX_SEPARATOR:
		state->s_current_arg->arg_type = ARG_INDEX_CONST;
		state->s_current_arg->arg_index_const.index_offset = tok;
		state->s_prev_component = INSTR_OPERAND_INDEX_OFFSET;
		return IVY_OK;
	default:
		return IVY_ERR_BAD_SYNTAX;
	}
}

static enum ivy_status parse_ident(
	struct ivy_asm_parser *ctx, struct ivy_asm_token *tok)
{
	struct block_parser_state *state
		= (struct block_parser_state *)asm_parser_get_state(ctx);

	unsigned long long x = 0;

	switch (state->s_prev_component) {
	case INSTR_NONE:
	case INSTR_OPCODE_DOT:
		b_queue_push_back(&state->s_mnemonic, &tok->t_entry);
		state->s_prev_component = INSTR_OPCODE;
		return IVY_OK;
	case INSTR_OPCODE:
	case INSTR_OPERAND_SEPARATOR:
		x = get_register_index(tok);
		state->s_prev_component = INSTR_OPERAND;

		if (x == REG_INDEX_INVALID) {
			return IVY_ERR_BAD_SYNTAX;
		} else {
			return push_reg_arg(state, tok, x);
		}
	case INSTR_OPERAND_INDEX_LEFT:
		x = get_index_base(tok);
		if (x == INDEX_NONE) {
			return IVY_ERR_BAD_SYNTAX;
		}

		state->s_current_arg->arg_type = ARG_INDEX_REG;
		state->s_current_arg->arg_index_reg.index_base_token = tok;
		state->s_current_arg->arg_index_reg.index_base = x;
		state->s_prev_component = INSTR_OPERAND_INDEX_BASE;
		return IVY_OK;
	default:
		return IVY_ERR_BAD_SYNTAX;
	}

	/* not sure what this is but we aren't expecting it. */
	return IVY_ERR_BAD_SYNTAX;
}

static enum ivy_status parse_label_ref(
	struct ivy_asm_parser *ctx, struct ivy_asm_token *tok)
{
	struct block_parser_state *state
		= (struct block_parser_state *)asm_parser_get_state(ctx);

	unsigned long long x = 0;

	switch (state->s_prev_component) {
	case INSTR_OPCODE:
	case INSTR_OPERAND_SEPARATOR:
		state->s_prev_component = INSTR_OPERAND;

		return push_label_arg(state, tok, x);
	default:
		return IVY_ERR_BAD_SYNTAX;
	}

	/* not sure what this is but we aren't expecting it. */
	return IVY_ERR_BAD_SYNTAX;
}

static enum ivy_status parse_dot(
	struct ivy_asm_parser *ctx, struct ivy_asm_token *tok)
{
	struct block_parser_state *state
		= (struct block_parser_state *)asm_parser_get_state(ctx);

	if (state->s_prev_component != INSTR_OPCODE) {
		return IVY_ERR_BAD_SYNTAX;
	}

	state->s_prev_component = INSTR_OPCODE_DOT;
	return IVY_OK;
}

static enum ivy_status parse_label(
	struct ivy_asm_parser *ctx, struct ivy_asm_token *tok)
{
	struct block_parser_state *state
		= (struct block_parser_state *)asm_parser_get_state(ctx);

	if (state->s_prev_component != INSTR_NONE) {
		/* not sure what this is but we aren't expecting it. */
		return IVY_ERR_BAD_SYNTAX;
	}

	return push_label(ctx, state, tok);
}

static enum ivy_status parse_comma(
	struct ivy_asm_parser *ctx, struct ivy_asm_token *tok)
{
	struct block_parser_state *state
		= (struct block_parser_state *)asm_parser_get_state(ctx);

	switch (state->s_prev_component) {
	case INSTR_OPERAND:
	case INSTR_OPERAND_INDEX_RIGHT:
		state->s_prev_component = INSTR_OPERAND_SEPARATOR;
		return IVY_OK;
	case INSTR_OPERAND_INDEX_BASE:
		state->s_prev_component = INSTR_OPERAND_INDEX_SEPARATOR;
		return IVY_OK;
	default:
		return IVY_ERR_BAD_SYNTAX;
	}
}

static enum ivy_status parse_left_bracket(
	struct ivy_asm_parser *ctx, struct ivy_asm_token *tok)
{
	struct block_parser_state *state
		= (struct block_parser_state *)asm_parser_get_state(ctx);

	switch (state->s_prev_component) {
	case INSTR_OPCODE:
	case INSTR_OPERAND_SEPARATOR:
		break;
	default:
		return IVY_ERR_BAD_SYNTAX;
	}

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

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

	state->s_current_arg = arg;
	state->s_prev_component = INSTR_OPERAND_INDEX_LEFT;
	return IVY_OK;
}

static enum ivy_status parse_right_bracket(
	struct ivy_asm_parser *ctx, struct ivy_asm_token *tok)
{
	struct block_parser_state *state
		= (struct block_parser_state *)asm_parser_get_state(ctx);

	if (state->s_prev_component != INSTR_OPERAND_INDEX_OFFSET) {
		return IVY_ERR_BAD_SYNTAX;
	}

	b_queue_push_back(&state->s_args, &state->s_current_arg->arg_entry);
	state->s_current_arg = NULL;
	state->s_prev_component = INSTR_OPERAND;

	return IVY_OK;
}

static enum ivy_status parse_end(
	struct ivy_asm_parser *ctx, struct ivy_asm_token *tok)
{
	struct block_parser_state *state
		= (struct block_parser_state *)asm_parser_get_state(ctx);

	if (state->s_prev_component != INSTR_NONE) {
		return IVY_ERR_BAD_SYNTAX;
	}

	asm_parser_pop_state(ctx, NULL);

	return IVY_OK;
}

static void init_state(struct ivy_asm_parser *ctx, struct parser_state *state)
{
	ivy_assembler_begin_scope(
		ctx->p_assembler, IVY_ASM_SCOPE_BLOCK, state->s_attrib);
}

static void finish_state(struct ivy_asm_parser *ctx, struct parser_state *state)
{
	ivy_assembler_end_scope(ctx->p_assembler);
}

struct parser_state_type block_parser_state_type = {
	.n_init_state = init_state,
	.n_finish_state = finish_state,
	.n_state_size = sizeof(struct block_parser_state),
	.n_token_parsers = {
		TOK_PARSER(IDENT, parse_ident),
		TOK_PARSER(LABEL, parse_label),
		TOK_PARSER(LABEL_REF, parse_label_ref),
		TOK_PARSER(INT, parse_int),
		TOK_PARSER(LINEFEED, parse_linefeed),
	},
	.n_symbol_parsers = {
		SYM_PARSER(LEFT_BRACKET, parse_left_bracket),
		SYM_PARSER(RIGHT_BRACKET, parse_right_bracket),
		SYM_PARSER(COMMA, parse_comma),
		SYM_PARSER(DOT, parse_dot),
	},
	.n_keyword_parsers = {
		KW_PARSER(END, parse_end),
	},
};