#include "ops.h"
#include "drain.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>

typedef struct Label Label;
struct Label {
	const char *s;
	size_t addr;
	Label *next;
};

static const int hextable[256] = {
	['0'] = 0x0, ['1'] = 0x1, ['2'] = 0x2, ['3'] = 0x3, ['4'] = 0x4,
	['5'] = 0x5, ['6'] = 0x6, ['7'] = 0x7, ['8'] = 0x8, ['9'] = 0x9,
	['a'] = 0xa, ['b'] = 0xb, ['c'] = 0xc, ['d'] = 0xd, ['e'] = 0xe,
	['f'] = 0xf,
};
static const char *sep = "\t\n\v\f\r ";
static int verbose = 0;
static char *data = NULL, *data_dup = NULL;
FILE *outfile = NULL;
static Label *label_root = NULL;

static int
fail(void)
{
	if (data != NULL)
		free(data);
	if (data_dup != NULL)
		free(data_dup);
	while (label_root != NULL) {
		Label *const next = label_root->next;
		free(label_root);
		label_root = next;
	}
	if (outfile != NULL && outfile != stdout)
		fclose(outfile);
	return 1;
}

static void
skip_comments(void)
{
	char *tok;
	do
		tok = strtok(NULL, sep);
	while (tok != NULL && strcmp(tok, ")") != 0);

	if (tok == NULL) {
		fprintf(stderr, "unclosed comment");
		exit(fail());
	}
}

static void
write_short(uint_least16_t v)
{
	static unsigned int write_n = 0;
	if (verbose) fprintf(stderr, "write %04x: %04x\n", write_n, v);
	write_n += 1;
	if (fputc(v & 0x00ff, outfile) == EOF ||
	    fputc(v / 0x0100, outfile) == EOF) {
		fprintf(stderr, "fputc failed\n");
		exit(fail());
	}
}

static void
register_label(const char *s, size_t pc)
{
	Label *const label = malloc(sizeof(Label));
	if (label == NULL) {
		perror("register_label:malloc");
		exit(fail());
	}

	label->s = s;
	label->addr = pc;
	label->next = label_root;
	label_root = label;
	if (verbose) fprintf(stderr, "label: '%s' at pc %zu\n", s, pc);
}

static void
write_label(const char *s)
{
	for (Label *label = label_root; label != NULL; label = label->next) {
		if (strcmp(s, label->s) == 0) {
			write_short(label->addr);
			return;
		}
	}
	fprintf(stderr, "unknown label '%s'\n", s);
}

static int
opcmp(const char *s0, const char *s1)
{
	return (s0[0] == s1[0] && s0[1] == s1[1] && s0[2] == s1[2]);
}

static void
write_opcode(const char *s)
{
	for (unsigned i = 0; i < sizeof(ops) / sizeof(ops[0]); i++) {
		if (opcmp(s, ops[i])) {
			write_short(i);
			return;
		}
	}
	fprintf(stderr, "unknown opcode '%s'\n", s);
	exit(fail());
}

static int
ishex(int c)
{
	return ((c >= 'a' && c <= 'f') || (c >= '0' && c <= '9'));
}

static void
write_literal(const char *s)
{
	if (!ishex(s[0]) && !ishex(s[1]) && !ishex(s[2]) && !ishex(s[3])) {
		fprintf(stderr, "invalid literal '%s'\n", s);
		exit(fail());
	}
	write_short((hextable[(int)s[0]] * 0x1000) |
	            (hextable[(int)s[1]] * 0x100) |
	            (hextable[(int)s[2]] * 0x10) |
	            (hextable[(int)s[3]]));
}

static void
first_pass(char *s)
{
	char *tok;
	size_t pc = 0;
	for (tok = strtok(s, sep); tok != NULL; tok = strtok(NULL, sep)) {
		/* Skip comment. */
		if (strcmp(tok, "(") == 0) {
			skip_comments();
			continue;
		}

		switch (tok[0]) {
		case '@': /* label marker */
			register_label(tok + 1, pc);
			break;
		case '/': /* call label */
		case '-': /* literal label */
		case '#': /* literal marker */
			pc += 2;
			break;
		default:
			pc += 1;
			break;
		}
	}
}

static void
second_pass(char *s)
{
	char *tok;
	for (tok = strtok(s, sep); tok != NULL; tok = strtok(NULL, sep)) {
		switch (tok[0]) {
		case '@': /* label declaration, skip */
			continue;
		case '/': /* call label */
			write_short(OP_JRT);
			write_label(tok + 1);
			continue;
		case '-': /* literal label insertion */
			write_short(OP_LIT);
			write_label(tok + 1);
			continue;
		case ',': /* label insertion */
			write_label(tok + 1);
			continue;
		default:
			break;
		}

		/* Skip comment. */
		if (strcmp(tok, "(") == 0) {
			skip_comments();
			continue;
		}

		const size_t tok_len = strlen(tok);
		switch (tok_len) {
		case 3: /* opcode */
			if (verbose) fprintf(stderr, "opcode: %s\n", tok);
			write_opcode(tok);
			break;
		case 5: /* LIT literal */
			if (tok[0] != '#')
				fprintf(stderr, "unknown token '%s'\n", tok);
			write_short(OP_LIT);
			tok += 1;
			/* fallthrough */
		case 4: /* literal */
			if (verbose) fprintf(stderr, "literal: %s\n", tok);
			write_literal(tok);
			break;
		default: /* wtf */
			fprintf(stderr, "unknown token '%s'\n", tok);
			exit(fail());
		}
	}
}

int
main(int argc, char **argv)
{
	if (!(argc == 3 || (argc == 4 && strcmp(argv[3], "-v") == 0))) {
		fprintf(stderr, "usage: %s <asm> <rom> [-v]\n", argv[0]);
		return fail();
	}
	verbose = (argc == 4);

	if (strcmp(argv[1], "-") == 0)
		data = drain_stdin(NULL);
	else {
		/* Dump file to memory. */
		FILE *const fp = fopen(argv[1], "rb");
		if (fp == NULL) {
			perror(argv[1]);
			return fail();
		}
		data = drain(fp, NULL);
		fclose(fp);
	}
	if (data == NULL)
		return fail();

	/* Duplicate data. */
	data_dup = malloc(strlen(data) + 1);
	if (data_dup == NULL) {
		perror("main:malloc");
		return fail();
	}
	strcpy(data_dup, data);

	/* Manually register standard addresses. */
	register_label("Input", 0xbffe);
	register_label("Screen", 0xbfff);

	/* First pass, find labels. */
	first_pass(data);

	if (strcmp(argv[2], "-") == 0)
		outfile = stdout;
	else {
		/* Open out file. */
		outfile = fopen(argv[2], "wb");
		if (outfile == NULL) {
			perror(argv[2]);
			return fail();
		}
	}

	/* Second pass, generate bytecode. */
	second_pass(data_dup);

	return fail(), 0;
}