#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#ifdef _DEBUG
#define DBGLOG   printf
#else
#define DBGLOG(...)
#endif

#define err(...) (fprintf(stderr, __VA_ARGS__), \
                  fprintf(stderr, "%lu | %s\n", file_line_no, file_line), false)
#define die(...)   do { fprintf(stderr, __VA_ARGS__); exit(-1); } while (0)

#define is_symbol_char(c) (('0' <= c && c <= '9') || ('A' <= c && c <= 'Z') \
                        || ('a' <= c && c <= 'z') || c == '_' || c == '.'   \
                                                  || c == '$' || c == ':')
#define is_whitespace(c) ((c == ' ' || c == '\t' || c == '\n' || c == '\r'))
#define is_number(c) (('0' <= c && c <= '9'))

#define MAX_LINE_LEN          256 // TODO: in/excludes NULL terminator?
#define MAX_SYMBOL_LEN        MAX_LINE_LEN - 2 + 1 // + 1 for NULL terminator

size_t instruction_offset = 0;

uint16_t rom[32768];
size_t rom_index = 0;

char  *file_line;     // reference to currently-read line (for convenience)
size_t file_line_no;  // line number, regardless of line content

// label and variable symbols are treated the same internally
#define MAX_SYMBOLS     32768
char *reserved_strs[] = {"R0", "R1", "R2", "R3", "R4", "R5", "R6", "R7",
                         "R8", "R9", "R10", "R11", "R12", "R13", "R14", "R15",
                         "SP", "LCL", "ARG", "THIS", "THAT", "SCREEN", "KBD"};
uint16_t reserved_vals[] = {0, 1, 2, 3, 4, 5, 6, 7,
                            8, 9, 10, 11, 12, 13, 14, 15,
                            0, 1, 2, 3, 4, 0x4000, 0x6000};
char     symbol_strs[MAX_SYMBOLS + sizeof(reserved_vals)][MAX_SYMBOL_LEN];
uint16_t symbol_vals[MAX_SYMBOLS + sizeof(reserved_vals)];
uint16_t symbol_index;
uint16_t next_var_address = 16;


static void write_rom_binary(FILE *fp)
{
	size_t i, k;
        char bin_str[17];
	for (i = 0; i < rom_index; ++i) {
		for (k = 0; k < 16; ++k, rom[i] <<= 1)
			bin_str[k] = (((rom[i] >> 15) & 0x01) != 0) ? '1' : '0';
		bin_str[16] = '\0';
		fprintf(fp, "%s\n", bin_str);
	}
}

static size_t skip_whitespace(const char *line, size_t n)
{
	size_t i;
	for (i = 0; is_whitespace(line[i]) && i < n; ++i);
	return i;
}

//static void debug_dump_all_symbols()
//{
//	size_t i;
//	DBGLOG("-------- DEBUG SYMBOL DUMP --------\n");
//	for (i = 0; i < symbol_index; ++i)
//		DBGLOG("symbol (%s, %hu)\n", symbol_strs[i], symbol_vals[i]);
//	DBGLOG("-------- END SYMBOL DUMP --------\n");
//}

static char *lookup_symbol(const char *str, uint16_t *val)
{
	size_t i;

	for (i = 0; i < sizeof(symbol_strs) / sizeof(symbol_strs[0]); ++i) {
		if (strncmp(symbol_strs[i], str, MAX_SYMBOL_LEN) == 0) {
			*val = symbol_vals[i];
			return symbol_strs[i];
		}
	}

	return NULL;
}

static bool add_symbol(const char *str, uint16_t val)
{
	uint16_t tmp;

	if (symbol_index > 32767)
		return err("error: failed to add symbol '%s': %hu > 32767\n",
		           str, val);

	if (lookup_symbol(str, &tmp) != NULL)
		return err("error: failed to add symbol '%s': symbol already "
		           "exists\n", str);

	strncpy(symbol_strs[symbol_index], str, MAX_SYMBOL_LEN);
	symbol_vals[symbol_index] = val;
	++symbol_index;

	return true;
}

// pre-fill symbol lists with 'reserved' symbols and values
static bool init_reserved_symbols()
{
	size_t i;

	for (i = 0; i < sizeof(reserved_strs) / sizeof(reserved_strs[0]); ++i)
		if (!add_symbol(reserved_strs[i], reserved_vals[i]))
			return err("error: failed to init reserved symbols\n");

	return true;
}

// assumes line[0] == '('
bool parse_line_for_label(const char *line) {
	size_t i;
	char c, label_str[MAX_SYMBOL_LEN + 1];

	for (i = 1; i < MAX_SYMBOL_LEN; ++i) {
		c = line[i];
		if (c == '\0' || c == ')')
			break;
		if (!is_symbol_char(c))
			return err("error: invalid char '%c' in label\n", c);
		label_str[i - 1] = line[i];
	}

	if (c != ')')
		return err("error: no matching ')' found for label\n");

	label_str[i - 1] = '\0';
	return add_symbol(label_str, (uint16_t)instruction_offset);
}

/* returns dest bits 0b00000ddd
 */
bool parse_c_type_dest(const char *line, uint16_t *dest)
{
	size_t i;
	uint8_t a_dest = 0, d_dest = 0, m_dest = 0;

	for (i = 0; line[i] != '='; ++i) {
		switch (line[i]) {
		case 'A': ++a_dest; break;
		case 'D': ++d_dest; break;
		case 'M': ++m_dest; break;
		default: return  err("syntax error: invalid destination "
			               "register '%c'\n", line[i]);
		}

		if (i >= 3)
			return err("error: too many destinations\n");
	}

	if (a_dest > 1)
		return err("syntax error: multiple 'A's in destination\n");
	if (d_dest > 1)
		return err("syntax error: multiple 'D's in destination\n");
	if (m_dest > 1)
		return err("syntax error: multiple 'M's in destination\n");

	*dest = (uint16_t)((a_dest << 2) | (d_dest << 1) | (m_dest));
	return true;
}

/* returns comp bits 0b0acccccc
 */
bool parse_c_type_comp(const char *line, uint16_t *comp)
{
	size_t i;
	char comp_str[5];
	char *comp_tb[] = {"0", "1", "-1", "D", "A", "M", "!D", "!A",
	                   "!M", "-D", "-A", "-M", "D+1", "A+1", "M+1", "D-1",
	                   "A-1", "M-1", "D+A", "D+M", "D-A", "D-M", "A-D",
	                   "M-D", "D&A", "D&M", "D|A", "D|M"};
	char comp_vals[] = {0x2a, 0x3f, 0x3a, 0x0c, 0x30, 0x70, 0x0d, 0x31,
	                     0x71, 0x0f, 0x33, 0x73, 0x1f, 0x37, 0x77, 0x0e,
	                     0x32, 0x72, 0x02, 0x42, 0x13, 0x53, 0x07,
	                     0x47, 0x00, 0x40, 0x15, 0x55};

	for (i = 0; i < 5 && !is_whitespace(line[i]) && line[i] != ';'
	                  && line[i] != '/' && line[i] != '\0'; ++i)
		comp_str[i] = line[i];

	comp_str[i] = '\0';
	if (i > 4)
		return err("syntax error: invalid comp field '%s'\n", comp_str);

	for (i = 0; i < (sizeof(comp_tb) / sizeof(comp_tb[0])); ++i) {
		if (strncmp(comp_tb[i], comp_str, 3) == 0) {
			*comp = (uint16_t)comp_vals[i];
			return true;
		}
	}

	return err("syntax error: invalid comp field '%s'\n", comp_str);
}

/* returns jump bits 0b00000jjj
 */
bool parse_c_type_jump(const char *line, uint16_t *jump)
{
	size_t i;
	char jump_str[5];
	char *jump_tb[] = {"", "JGT", "JEQ", "JGE", "JLT", "JNE", "JLE", "JMP"};

	for (i = 0; i < 5 && !is_whitespace(line[i]) && line[i] != '/'
	                  && line[i] != '\0'; ++i)
		jump_str[i] = line[i];

	jump_str[i] = '\0';
	if (i > 4)
		return err("syntax error: invalid jump field '%s'\n", jump_str);

	for (i = 0; i < (sizeof(jump_tb) / sizeof(jump_tb[0])); ++i) {
		if (strncmp(jump_tb[i], jump_str, 3) == 0) {
			*jump = (uint16_t)i;
			return true;
		}
	}

	return err("syntax error2: invalid jump field '%s'\n", jump_str);
}

/* instruction format: 0b111accccccdddjjj
 */
bool parse_c_type(const char *line, uint16_t *instruction)
{
	size_t i;
	char c;
	const char *dest_start = NULL, *comp_start = NULL, *jump_start = NULL;
	uint16_t dest = 0, comp = 0, jump = 0;

	c = line[0];
	for (i = 0; !is_whitespace(c) && c != '\0'; ++i) {
		c = line[i];
		if (c == '=') {
			dest_start = &line[0];
			comp_start = &line[i+1];
		} else if (c == ';')
			jump_start = &line[i+1];
	}

	// only comp field is mandatory; dest/jump are optional, may not exist
	if (comp_start == NULL)
		comp_start = &line[0];

	if (!parse_c_type_comp(comp_start, &comp))
		return false;

	if (dest_start != NULL)
		if (!parse_c_type_dest(dest_start, &dest))
			return false;

	if (jump_start != NULL)
		if (!parse_c_type_jump(jump_start, &jump))
			return false;

	*instruction = 0xe000 | (comp << 6) | (dest << 3) | jump;
	return true;
}

bool parse_a_type(const char *line, uint16_t *instruction)
{
	char a_str[MAX_SYMBOL_LEN + 1], *sym_ptr;
	size_t i, a_str_len;
	uint16_t a_field = 0;

	for (i = 1; i < MAX_SYMBOL_LEN && is_symbol_char(line[i]); ++i) {
		a_str[i - 1] = line[i];
	}

	if (!(is_whitespace(line[i]) || line[i] == '\0' || line[i] == '/'))
		return err("syntax error: invalid char '%c' in instruction\n",
		           line[i]);

	a_str_len = i - 1;
	a_str[a_str_len] = '\0';

	if (a_str_len == 0)
		return err("syntax error: instruction empty after @\n");

	if (is_number(a_str[0])) {  // @<string> probably number
		if (a_str_len > 5)
			return err("syntax error: instruction field '%s' too "
			      "large\n", a_str);
		for (i = 0; i < 5 && is_number(a_str[i]); ++i) {
			if (!is_number(a_str[i]))
				return err("syntax error: '%s' not a valid "
				           "number\n", a_str);
			a_field = (a_field * 10) + (a_str[i] - 0x30);
		}
		if (a_field > 0x7fff)
			return err("syntax error: %u > 32767, too large\n",
			           a_field);
	} else {  // @<string> probably symbol
		sym_ptr = lookup_symbol(a_str, &a_field);
		if (sym_ptr == NULL) {
			a_field = next_var_address++;
			if (!add_symbol(a_str, a_field))
				return err("error: failed to add var '%s'\n",
				           a_str);
		}
	}

	*instruction = a_field;
	return true;
}

// does not care about line length; exits at first newline or after relevant
// portion parsed (allows for syntactically-incorrect lines, I know)
bool parse_instruction(const char *line, uint16_t *instruction)
{
	if (line[0] == '@')
		return parse_a_type(line, instruction);
	else if (line[0] >= '!' && line[0] < '~')
		return parse_c_type(line, instruction);

	return err("syntax error: invalid char '%c' found\n", line[0]);
}

// if first_pass == true: read labels and associate them with values
// else: (second pass) parse instructions
// Still unsure how I feel about doing handling both passes in one function.
// It's hacky, but also space saving and elegant in a way.
bool pass(FILE *in_file, bool first_pass)
{
	bool ret = true, second_pass = !first_pass; 
	char *line, in_line[MAX_LINE_LEN];
	uint16_t instruction;
	size_t i, line_len;

	if (first_pass)
		if (!init_reserved_symbols())
			return false;

	file_line_no = 0;
	instruction_offset = 0;
	while (fgets(in_line, MAX_LINE_LEN, in_file) != NULL) {  // parse loop
		++file_line_no;
		//DBGLOG("line %lu: %s", file_line_no, in_line);
		file_line = in_line;
		line = &in_line[skip_whitespace(in_line, MAX_LINE_LEN)];

		line_len = 0;
		for (i = 0; line[i] != '\0'; ++i) {
			if (line[i] == '\n' || line[i] == '\r') {
				line[i] = '\0'; // remove newlines
				break;
			}
			++line_len;  // get line length
		}

		if (line_len == 0)  // "empty" line
			continue;
		if (line_len > 1)
			if (line[0] == '/' && line[1] == '/')  // comment found
				continue;
		if (line[0] == '(') {
			if (first_pass) {
				if (!parse_line_for_label(line))
					ret = false;
			} else
				continue;
		} else {
			if (second_pass) {
				if (parse_instruction(line, &instruction))
					rom[rom_index++] = instruction;
				else
					ret = false;
			}
			++instruction_offset;
			continue;  // first pass: if not label, ignore line
		}
	}

	return ret;
}


char *usage_msg = "Usage: n2tasm2 [input/file.asm] [assembled/output]\n"
                  "       n2tasm2 [input/file.asm] # will output to STDOUT\n";

int main(int argc, char *argv[])
{
	FILE *in_file, *out_file;
	char *in_file_path;
	if (argc == 2) {
		out_file = stdout;
	} else if (argc == 3) {
		out_file = fopen(argv[2], "wb");
		if (out_file == NULL)
			die("failed to open %s for writing\n", argv[2]);
	} else {
		die(usage_msg);
	}

	in_file_path = argv[1];
	in_file = fopen(in_file_path, "r");
	if (in_file == NULL)
		die("failed to open assembly file for reading\n");

	if(!pass(in_file, true)) // first pass
		die("failed to parse labels/variables in file\n");

	//debug_dump_all_symbols();
	if (fseek(in_file, 0, SEEK_SET))
		die("failed to re-read file from beginning\n");

	if(!pass(in_file, false)) // second pass
		die("failed to parse assembly in file\n");

	write_rom_binary(out_file);
	if (fclose(in_file))
		die("failed to close assembly file\n");
	if (fclose(out_file))
		die("failed to close output file\n");

	return 0;
}