#include "iccmp.h"
#include "allocator.h"
#include "aoc.h"
#include "hmap.h"
#include "maxint.h"
#include "util.h"

#include <assert.h>
#include <stdbool.h>
#include <stdint.h>

const char *icc_err[] = {
	[ICC_OK] = "ok",
	[ICC_INV_INST] = "invalid instruction",
	[ICC_OOB_WRITE] = "write addr oob",
	[ICC_OOB_READ] = "read addr oob",
};

static struct maxint *
mi_tmp(int64_t val)
{
	static mi_ul data = 0;
	static struct maxint m = { &data, 1, 1, MI_POS };

	data = (mi_ul) ABS(val);
	m.sign = val >= 0 ? MI_POS : MI_NEG;

	return &m;
}

static bool
icc_hmap_keycmp(struct hmap_key k1, struct hmap_key k2)
{
	return mi_cmp(k1.p, k2.p) == 0;
}

static uint32_t
icc_hmap_hash(struct hmap_key key)
{
	return (uint32_t) mi_cast_ul(key.p);
}

void
icc_init(struct icc *icc)
{
	int rc;

	icc->state = ICC_RUN;

	icc->debug = getenv("ICC_DEBUG") != NULL;

	mi_init(&icc->rip);
	mi_setv(&icc->rip, 0, MI_POS);

	mi_init(&icc->base);
	mi_setv(&icc->base, 0, MI_POS);

	mi_init(&icc->in);
	mi_init(&icc->out);
	mi_init(&icc->read_addr);
	mi_init(&icc->write_addr);

	mi_init(&icc->r1);
	mi_init(&icc->r2);
	mi_init(&icc->r3);
	mi_init(&icc->r4);
	mi_init(&icc->tmp);

	icc->abort_on_err = true;

	icc->line_terminated = true;

	rc = hmap_init(&icc->instructions, 1024, icc_hmap_hash,
		icc_hmap_keycmp, &stdlib_strict_heap_allocator);
	assert(!rc);
}

void
icc_deinit(struct icc *icc)
{
	struct hmap_iter iter;

	mi_deinit(&icc->rip);
	mi_deinit(&icc->base);
	mi_deinit(&icc->in);
	mi_deinit(&icc->out);
	mi_deinit(&icc->read_addr);
	mi_deinit(&icc->write_addr);
	mi_deinit(&icc->r1);
	mi_deinit(&icc->r2);
	mi_deinit(&icc->r3);
	mi_deinit(&icc->r4);
	mi_deinit(&icc->tmp);

	for (HMAP_ITER(&icc->instructions, iter)) {
		mi_deinit(iter.link->key._p);
		free(iter.link->key._p);
		mi_deinit(iter.link->value._p);
		free(iter.link->value._p);
	}

	hmap_deinit(&icc->instructions);
}

void
icc_copy(struct icc *dst, struct icc *src)
{
	struct hmap_iter hmap_iter;
	struct maxint *key, *value;

	dst->state = src->state;
	mi_set(&dst->rip, &src->rip);
	mi_set(&dst->base, &src->base);
	mi_set(&dst->in, &src->in);
	mi_set(&dst->out, &src->out);

	for (HMAP_ITER(&dst->instructions, hmap_iter)) {
		mi_deinit(hmap_iter.link->key._p);
		free(hmap_iter.link->key._p);
		mi_deinit(hmap_iter.link->value._p);
		free(hmap_iter.link->value._p);
	}
	hmap_clear(&dst->instructions);

	for (HMAP_ITER(&src->instructions, hmap_iter)) {
		key = malloc(sizeof(struct maxint));
		mi_init(key);
		mi_set(key, hmap_iter.link->key.p);
		value = malloc(sizeof(struct maxint));
		mi_init(value);
		mi_set(value, hmap_iter.link->value.p);
		hmap_add(&dst->instructions, (struct hmap_key) {.p = key},
			(struct hmap_val) {.p = value});
	}
}

void
icc_reset(struct icc *icc, struct hmap *inst)
{
	int rc;

	icc->state = ICC_RUN;
	mi_setv(&icc->rip, 0, MI_POS);
	mi_setv(&icc->base, 0, MI_POS);
	if (inst) {
		rc = hmap_copy(&icc->instructions, inst);
		assert(!rc);
	}
}

void
icc_parse_inst(struct icc *icc, const char *str, size_t len)
{
	struct maxint *key, *val;
	const char *pos, *end;
	char buf[256];
	int64_t v;
	int rc, rip;

	pos = str;
	end = str + len;
	rip = 0;
	while (readtok(buf, sizeof(buf), ',', &pos, end)) {
		v = parsei64(buf);

		key = malloc(sizeof(struct maxint));
		assert(key != NULL);
		mi_init(key);
		mi_setv(key, (mi_ul) rip, MI_POS);

		val = malloc(sizeof(struct maxint));
		assert(val != NULL);
		mi_init(val);
		mi_setv(val, (mi_ul) ABS(v), v >= 0 ? MI_POS : MI_NEG);

		rc = hmap_add(&icc->instructions, (struct hmap_key) {.p = key},
			(struct hmap_val) {.p = val});
		assert(!rc);

		rip += 1;
	}
}

const char *
icc_literal_str(struct icc *icc, struct maxint *addr, size_t zfill)
{
	static char buf[64];

	mi_dec(addr, buf, sizeof(buf), zfill);

	return buf;
}

const char *
icc_value_str(struct icc *icc, struct maxint *addr, size_t zfill)
{
	static char buf[62];
	struct maxint *val;
	int rc;

	rc = icc_read(icc, addr, &val);
	if (rc) {
		snprintf(buf, sizeof(buf), "???");
	} else {
		mi_dec(val, buf, sizeof(buf), zfill);
	}

	return buf;
}

static void
icc_debug_op_pre(struct icc *icc)
{
	if (!icc->debug) return;

	fprintf(stderr, "%*s: ", 5, icc_literal_str(icc, &icc->rip, 0));
	fprintf(stderr, "(%*s) ", 5, icc_value_str(icc, &icc->rip, 5));
	icc->line_terminated = false;
}

static void
icc_debug_op_main(struct icc *icc, struct maxint *tmp, const char *opstr, int n)
{
	struct maxint *inst, *val;
	int i, rc, instint;

	if (!icc->debug) return;

	rc = icc_read(icc, &icc->rip, &inst);
	assert(!rc);

	fprintf(stderr, "%s ", opstr);
	icc->line_terminated = false;

	mi_set(tmp, &icc->rip);
	instint = (int) mi_cast_ul(inst);
	for (i = 1; i <= n; i++) {
		if (i > 1) fprintf(stderr, ", ");
		mi_add(tmp, tmp, mi_tmp(1));
		rc = icc_read(icc, tmp, &val);
		if (rc) {
			fprintf(stderr, "???");
			continue;
		}
		switch (icc_param_mode(instint, i)) {
		case ICC_PARAM_IMM:
			fprintf(stderr, "%s", icc_literal_str(icc, val, 0));
			break;
		case ICC_PARAM_POS:
			fprintf(stderr, "[%s]=%s", icc_literal_str(icc, val, 0),
				icc_value_str(icc, val, 0));
			break;
		case ICC_PARAM_REL:
			mi_add(tmp, &icc->base, val);
			fprintf(stderr, "[%s+", icc_literal_str(icc, &icc->base, 0));
			fprintf(stderr, "%s", icc_literal_str(icc, val, 0));
			fprintf(stderr, "=%s]=%s", icc_literal_str(icc, tmp, 0),
				icc_value_str(icc, tmp, 0));
			break;
		default:
			assert(0);
		}
	}
}

static void
icc_debug_op_post(struct icc *icc, struct maxint *addr)
{
	if (!icc->debug) return;

	fprintf(stderr, " -> [%s]=%s\n", icc_literal_str(icc, addr, 0),
		icc_value_str(icc, addr, 0));
	icc->line_terminated = true;
}

static void
icc_debug_op(struct icc *icc, struct maxint *tmp, const char *opstr, int n)
{
	if (!icc->debug) return;

	icc_debug_op_main(icc, tmp, opstr, n);
	fprintf(stderr, "\n");
	icc->line_terminated = true;
}

int
icc_inst_add(struct icc *icc)
{
	struct maxint *dst;
	int rc;

	icc_debug_op_main(icc, &icc->tmp, "ADD", 3);

	rc = icc_get_param(icc, &icc->tmp, 1, &icc->r1);
	if (rc) return rc;
	rc = icc_get_param(icc, &icc->tmp, 2, &icc->r2);
	if (rc) return rc;
	rc = icc_get_dest(icc, &icc->tmp, 3, &icc->r3);
	if (rc) return rc;

	mi_add(&icc->tmp, &icc->r1, &icc->r2);
	rc = icc_write(icc, &icc->r3, &icc->tmp);
	if (rc) return rc;

	icc_debug_op_post(icc, dst);

	mi_add(&icc->rip, &icc->rip, mi_tmp(4));
	icc->state = ICC_RUN;

	return 0;
}

int
icc_inst_mul(struct icc *icc)
{
	int rc;

	icc_debug_op_main(icc, &icc->tmp, "MUL", 3);

	rc = icc_get_param(icc, &icc->tmp, 1, &icc->r1);
	if (rc) return rc;
	rc = icc_get_param(icc, &icc->tmp, 2, &icc->r2);
	if (rc) return rc;
	rc = icc_get_dest(icc, &icc->tmp, 3, &icc->r3);
	if (rc) return rc;

	mi_mul(&icc->tmp, &icc->r1, &icc->r2);
	rc = icc_write(icc, &icc->r3, &icc->tmp);
	if (rc) return rc;

	icc_debug_op_post(icc, &icc->r3);

	mi_add(&icc->rip, &icc->rip, mi_tmp(4));
	icc->state = ICC_RUN;

	return 0;
}

int
icc_inst_store(struct icc *icc)
{
	int rc;

	if (icc->state != ICC_INPUT || !icc->in.size) {
		icc_debug_op(icc, &icc->tmp, "INPUT", 0);
		icc->state = ICC_INPUT;
		icc->in.size = 0;
	} else {
		icc_debug_op_main(icc, &icc->tmp, "STORE", 1);

		rc = icc_get_dest(icc, &icc->tmp, 1, &icc->r1);
		if (rc) return rc;

		rc = icc_write(icc, &icc->r1, &icc->in);
		if (rc) return rc;

		icc_debug_op_post(icc, &icc->r1);

		mi_add(&icc->rip, &icc->rip, mi_tmp(2));
		icc->state = ICC_RUN;
	}

	return 0;
}

int
icc_inst_load(struct icc *icc)
{
	int rc;

	if (icc->state != ICC_OUTPUT) {
		icc_debug_op(icc, &icc->tmp, "LOAD", 1);

		rc = icc_get_param(icc, &icc->tmp, 1, &icc->r1);
		if (rc) return rc;

		mi_set(&icc->out, &icc->r1);
		icc->state = ICC_OUTPUT;

		icc_debug_op(icc, &icc->tmp, "OUTPUT", 0);
	} else {
		mi_add(&icc->rip, &icc->rip, mi_tmp(2));
		icc->state = ICC_RUN;
	}

	return 0;
}

int
icc_inst_jmp_true(struct icc *icc)
{
	int rc;

	icc_debug_op(icc, &icc->tmp, "JMPT", 2);

	rc = icc_get_param(icc, &icc->tmp, 1, &icc->r1);
	if (rc) return rc;
	rc = icc_get_param(icc, &icc->tmp, 2, &icc->r2);
	if (rc) return rc;

	if (!mi_zero(&icc->r1)) {
		mi_set(&icc->rip, &icc->r2);
	} else {
		mi_add(&icc->rip, &icc->rip, mi_tmp(3));
	}

	icc->state = ICC_RUN;

	return 0;
}

int
icc_inst_jmp_false(struct icc *icc)
{
	int rc;

	icc_debug_op(icc, &icc->tmp, "JMPF", 2);

	rc = icc_get_param(icc, &icc->tmp, 1, &icc->r1);
	if (rc) return rc;
	rc = icc_get_param(icc, &icc->tmp, 2, &icc->r2);
	if (rc) return rc;

	if (mi_zero(&icc->r1)) {
		mi_set(&icc->rip, &icc->r2);
	} else {
		mi_add(&icc->rip, &icc->rip, mi_tmp(3));
	}

	icc->state = ICC_RUN;

	return 0;
}

int
icc_inst_test_lt(struct icc *icc)
{
	int rc;

	icc_debug_op_main(icc, &icc->tmp, "TLT", 3);

	rc = icc_get_param(icc, &icc->tmp, 1, &icc->r1);
	if (rc) return rc;
	rc = icc_get_param(icc, &icc->tmp, 2, &icc->r2);
	if (rc) return rc;
	rc = icc_get_dest(icc, &icc->tmp, 3, &icc->r3);
	if (rc) return rc;

	mi_setv(&icc->tmp, mi_cmp(&icc->r1, &icc->r2) == -1, MI_POS);
	rc = icc_write(icc, &icc->r3, &icc->tmp);
	if (rc) return rc;

	icc_debug_op_post(icc, &icc->r3);

	mi_add(&icc->rip, &icc->rip, mi_tmp(4));
	icc->state = ICC_RUN;

	return 0;
}

int
icc_inst_test_eq(struct icc *icc)
{
	int rc;

	icc_debug_op_main(icc, &icc->tmp, "TEQ", 3);

	rc = icc_get_param(icc, &icc->tmp, 1, &icc->r1);
	if (rc) return rc;
	rc = icc_get_param(icc, &icc->tmp, 2, &icc->r2);
	if (rc) return rc;
	rc = icc_get_dest(icc, &icc->tmp, 3, &icc->r3);
	if (rc) return rc;

	mi_setv(&icc->tmp, mi_cmp(&icc->r1, &icc->r2) == 0, MI_POS);
	rc = icc_write(icc, &icc->r3, &icc->tmp);
	if (rc) return rc;

	icc_debug_op_post(icc, &icc->r3);

	mi_add(&icc->rip, &icc->rip, mi_tmp(4));
	icc->state = ICC_RUN;

	return 0;
}

int
icc_inst_base(struct icc *icc)
{
	int rc;

	icc_debug_op_main(icc, &icc->tmp, "BASE", 1);

	rc = icc_get_param(icc, &icc->tmp, 1, &icc->r1);
	if (rc) return rc;

	mi_add(&icc->base, &icc->base, &icc->r1);
	mi_add(&icc->rip, &icc->rip, mi_tmp(2));
	icc->state = ICC_RUN;

	if (icc->debug) {
		fprintf(stderr, " -> BASE:%s\n",
			icc_literal_str(icc, &icc->base, 0));
		icc->line_terminated = true;
	}

	return 0;
}

int
icc_inst_halt(struct icc *icc)
{
	icc_debug_op(icc, &icc->tmp, "HALT", 0);

	icc->state = ICC_HALT;

	return 0;
}

int
icc_step_inst(struct icc *icc)
{
	struct maxint *inst;
	int rc;

	icc_debug_op_pre(icc);

	rc = icc_read(icc, &icc->rip, &inst);
	if (rc) return rc;

	switch (mi_cast_ul(inst) % 100) {
	case ICC_INST_ADD:
		rc = icc_inst_add(icc);
		break;
	case ICC_INST_MULT:
		rc = icc_inst_mul(icc);
		break;
	case ICC_INST_STORE:
		rc = icc_inst_store(icc);
		break;
	case ICC_INST_LOAD:
		rc = icc_inst_load(icc);
		break;
	case ICC_INST_JMPT:
		rc = icc_inst_jmp_true(icc);
		break;
	case ICC_INST_JMPF:
		rc = icc_inst_jmp_false(icc);
		break;
	case ICC_INST_TLT:
		rc = icc_inst_test_lt(icc);
		break;
	case ICC_INST_TEQ:
		rc = icc_inst_test_eq(icc);
		break;
	case ICC_INST_BASE:
		rc = icc_inst_base(icc);
		break;
	case ICC_INST_HALT:
		rc = icc_inst_halt(icc);
		break;
	default:
		rc = ICC_INV_INST;
		break;
	}

	if (icc->debug) {
		if (!icc->line_terminated) {
			fprintf(stderr, "\n");
			icc->line_terminated = true;
		}
		if (rc != ICC_OK)
			fprintf(stderr, " -- error: %s --\n", icc_err[rc]);
	}

	return rc;
}

struct hmap_link *
icc_write_any(struct icc *icc, struct maxint *addr, struct maxint *value)
{
	struct hmap_link **link;
	struct maxint *key, *val;

	link = hmap_link_pos(&icc->instructions, (struct hmap_key) {.p = addr});
	if (*link) {
		mi_set((*link)->value._p, value);
	} else {
		key = malloc(sizeof(struct maxint));
		assert(key != NULL);
		mi_init(key);
		mi_set(key, addr);

		val = malloc(sizeof(struct maxint));
		assert(val != NULL);
		mi_init(val);
		mi_set(val, value);

		*link = hmap_link_alloc(&icc->instructions,
			(struct hmap_key) {.p = key},
			(struct hmap_val) {.p = val}, NULL);
	}

	return *link;
}

int
icc_write(struct icc *icc, struct maxint *addr, struct maxint *val)
{
	struct hmap_link *link;

	link = hmap_get(&icc->instructions, (struct hmap_key) {.p = addr});
	if (!link) {
		if (addr->sign != MI_POS) {
			mi_set(&icc->read_addr, addr);
			assert(!icc->abort_on_err);
			return ICC_OOB_READ;
		}

		icc_write_any(icc, addr, val);
	} else {
		mi_set(link->value._p, val);
	}

	return 0;
}

int
icc_read(struct icc *icc, struct maxint *addr, struct maxint **out)
{
	struct hmap_link *link;

	link = hmap_get(&icc->instructions, (struct hmap_key) {.p = addr});
	if (!link) {
		if (addr->sign != MI_POS) {
			mi_set(&icc->read_addr, addr);
			assert(!icc->abort_on_err);
			return ICC_OOB_READ;
		}

		link = icc_write_any(icc, addr, mi_tmp(0));
	}
	*out = link->value._p;

	return 0;
}

int
icc_param_mode(int inst, int param)
{
	int div, i;

	div = 100;
	for (i = 1; i < param; i++) div *= 10;

	return (inst / div) % 10;
}

int
icc_get_param(struct icc *icc, struct maxint *tmp, int param, struct maxint *out)
{
	struct maxint *inst, *val, *val2;
	int rc;

	/* TODO: make out double pointer and optimize out mi_set */

	rc = icc_read(icc, &icc->rip, &inst);
	if (rc) return rc;

	mi_setv(tmp, (mi_ul) ABS(param), param >= 0 ? MI_POS : MI_NEG);
	mi_add(tmp, tmp, &icc->rip);
	rc = icc_read(icc, tmp, &val);
	if (rc) return rc;

	assert(mi_lastset(inst) == 0);
	switch (icc_param_mode((int) mi_cast_ul(inst), param)) {
	case ICC_PARAM_IMM:
		mi_set(out, val);
		break;
	case ICC_PARAM_POS:
		rc = icc_read(icc, val, &val2);
		if (rc) return rc;
		mi_set(out, val2);
		break;
	case ICC_PARAM_REL:
		mi_add(tmp, &icc->base, val);
		rc = icc_read(icc, tmp, &val2);
		mi_set(out, val2);
		break;
	default:
		assert(0);
	};

	return 0;
}

int
icc_get_dest(struct icc *icc, struct maxint *tmp, int param, struct maxint *out)
{
	struct maxint *inst, *val;
	int rc;

	rc = icc_read(icc, &icc->rip, &inst);
	if (rc) return rc;

	mi_setv(tmp, (mi_ul) ABS(param), param >= 0 ? MI_POS : MI_NEG);
	mi_add(tmp, tmp, &icc->rip);
	rc = icc_read(icc, tmp, &val);
	if (rc) return rc;

	switch (icc_param_mode((int) mi_cast_ul(inst), param)) {
	case ICC_PARAM_POS:
		mi_set(out, val);
		break;
	case ICC_PARAM_REL:
		mi_add(out, &icc->base, val);
		break;
	default:
		assert(0);
	}

	return 0;
}

void
icc_debug_print(struct icc *icc, struct maxint *addr)
{
	if (!icc->debug) return;

	fprintf(stderr, "[%s] = %s\n", icc_literal_str(icc, addr, 0),
		icc_value_str(icc, addr, 0));
}