#include "tpu.h"
#include "util.h"

#include <ctype.h>
#include <assert.h>
#include <stdio.h>
#include <stdbool.h>
#include <stddef.h>
#include <string.h>

const char *dir_reprs[] = {
	"LEFT", "RIGHT", "UP", "DOWN"
};

const char *mode_reprs[] = {
	"IDLE", "RUN", "READ", "WRITE"
};

const char *inst_reprs[] = {
	"NOP", "MOV", "SWP", "SAV", "ADD", "SUB",
	"NEG", "JMP", "JEZ", "JNZ", "JGZ", "JLZ", "JRO"
};

const char *op_reprs[] = {
	"ACC", "NIL", "LEFT", "RIGHT", "UP", "DOWN",
	"ANY", "LAST", "LIT", "<NAME>"
};

static enum tpu_port_dir
op_to_dir(enum tpu_inst_op_type op)
{
	return op - OP_LEFT;
}

static enum tpu_port_dir
opposite_dir(enum tpu_port_dir dir)
{
	switch (dir) {
	case DIR_UP:
		return DIR_DOWN;
	case DIR_DOWN:
		return DIR_UP;
	case DIR_LEFT:
		return DIR_RIGHT;
	case DIR_RIGHT:
		return DIR_LEFT;
	}
}

static uint32_t
djb_hash(const char *str)
{
	uint32_t hash;
	const char *c;

	hash = 5381;
	for (c = str; *c; c++)
		hash = 33 * hash + (uint8_t) tolower(*c);

	return hash;
}

void
label_map_init(struct label_map *map)
{
	memset(map, 0, sizeof(struct label_map));
}

void
label_map_deinit(struct label_map *map)
{
	struct label_map_link *link, *next;
	size_t i;

	for (i = 0; i < LABEL_MAP_BUCKETS; i++) {
		link = map->buckets[i];
		while (link) {
			next = link->next;
			free(link);
			link = next;
		}
	}

	for (i = 0; i < TPU_MAX_INST_CNT; i++)
		free(map->labels[i].str);
}

static struct label_map_link **
label_map_link_pos(struct label_map *map, const char *name)
{
	struct label_map_link **link;

	link = &map->buckets[djb_hash(name) % LABEL_MAP_BUCKETS];
	while (*link && strcasecmp(map->labels[(*link)->pc].str, name))
		link = &(*link)->next;

	return link;
}

bool
label_map_add(struct label_map *map, const char *name, int pc, bool prefix)
{
	struct label_map_link **pos, *link;

	pos = label_map_link_pos(map, name);
	if (*pos) return false;

	if (map->labels[pc].str) return false;
	map->labels[pc].str = strdup(name);
	if (!map->labels[pc].str) die("strdup:");
	map->labels[pc].prefix = prefix;

	*pos = link = malloc(sizeof(struct label_map_link));
	if (!link) die("malloc:");
	link->pc = pc;
	link->next = NULL;

	return true;
}

int
label_map_get(struct label_map *map, const char *name)
{
	struct label_map_link **link;

	link = label_map_link_pos(map, name);
	if (!*link) return -1;

	return (*link)->pc;
}

void
tpu_port_init(struct tpu_port *port, struct tpu *tpu)
{
	port->tpu = tpu;
	port->attached = false;
	port->dst_port = NULL;
	port->type = PORT_BIDI;
	port->reset_in = false;
	port->writing = false;
	port->avail = false;
	port->in = -1;
	port->writing = false;
	port->out = -1;
	port->io = false;
}

void
tpu_port_deinit(struct tpu_port *port)
{
	/* empty */
}

void
tpu_port_update(struct tpu_port *port)
{
	struct tpu *tpu;

	if (!port->attached) return;

	if (port->writing && !port->dst_port->avail) {
		port->dst_port->in = port->out;
		port->dst_port->avail = true;
	}

	if (port->reset_in && (!port->tpu || !port->tpu->idle)) {
		port->avail = false;
		port->dst_port->writing = false;
		if ((tpu = port->dst_port->tpu)) {
			/* hacky: switch mode to RUN after receive */
			assert(tpu->mode == MODE_WRITE);
			tpu->mode = MODE_RUN;
			tpu->idle = false;
			if (++tpu->pc >= tpu->inst_cnt) tpu->pc = 0;
		}
		port->reset_in = false;
	}
}

void
tpu_io_port_init(struct tpu_io_port *io_port, char c, enum tpu_port_dir dir,
	enum tpu_port_type type, int x, int y)
{
	tpu_port_init(&io_port->port, NULL);
	io_port->port.io = io_port;
	io_port->io_step = 0;
	io_port->type = type;
	io_port->file = NULL;
	io_port->dir = dir;
	io_port->c = c;
	io_port->x = x;
	io_port->y = y;
}

void
tpu_io_port_deinit(struct tpu_io_port *io_port)
{
	tpu_port_deinit(&io_port->port);
	if (io_port->file) fclose(io_port->file);
}

void
tpu_init(struct tpu *tpu)
{
	size_t i;

	tpu->disabled = false;

	tpu->mode = MODE_IDLE;
	tpu->idle = false;

	tpu->x = 0;
	tpu->y = 0;

	tpu->steps = 0;
	tpu->idle_steps = 0;

	tpu->pc = 0;
	tpu->acc = 0;
	tpu->bak = 0;
	tpu->inst_cnt = 0;
	tpu->rows = 0;
	label_map_init(&tpu->label_map);

	tpu->last = -1;
	tpu->io_port = -1;
	for (i = 0; i < 4; i++)
		tpu_port_init(&tpu->ports[i], tpu);
}

void
tpu_deinit(struct tpu *tpu)
{
	int i;

	for (i = 0; i < tpu->inst_cnt; i++) {
		if (tpu->insts[i].opcnt >= 1
				&& tpu->insts[i].ops[0].type == OP_LABEL)
			free(tpu->insts[i].ops[0].val.label);
	}
	label_map_deinit(&tpu->label_map);
}

struct tpu_inst *
tpu_current_inst(struct tpu *tpu)
{
	if (tpu->pc >= tpu->inst_cnt)
		return NULL;
	return &tpu->insts[tpu->pc];
}

void
tpu_init_ports(struct tpu *tpu, struct tpu_map *map)
{
	struct tpu *neighbor;
	enum tpu_port_dir odir;
	int x, y, i;

	for (i = 0; i < 4; i++) {
		if (tpu->ports[i].attached) continue;
		if (tpu->disabled) continue;

		switch (i) {
		case DIR_LEFT:
			x = tpu->x - 1;
			y = tpu->y;
			break;
		case DIR_RIGHT:
			x = tpu->x + 1;
			y = tpu->y;
			break;
		case DIR_UP:
			x = tpu->x;
			y = tpu->y - 1;
			break;
		case DIR_DOWN:
			x = tpu->x;
			y = tpu->y + 1;
			break;
		}

		neighbor = tpu_map_get(map, x, y);
		if (neighbor && !neighbor->disabled) {
			odir = opposite_dir((enum tpu_port_dir) i);
			tpu->ports[i].attached = true;
			tpu->ports[i].dst_port = &neighbor->ports[odir];
			neighbor->ports[odir].attached = true;
			neighbor->ports[odir].dst_port = &tpu->ports[i];
		}
	}
}

void
tpu_update(struct tpu *tpu)
{
	struct tpu_port *port;
	int i;

	for (i = 0; i < 4; i++) {
		port = &tpu->ports[i];
		tpu_port_update(port);
	}
}

bool
tpu_set_inst(struct tpu *tpu, int pc, enum tpu_inst_type inst_type,
	unsigned opcnt, struct tpu_inst_op op1, struct tpu_inst_op op2)
{
	struct tpu_inst *inst;

	inst = &tpu->insts[pc];
	inst->type = inst_type;
	inst->opcnt = opcnt;
	inst->ops[0] = op1;
	inst->ops[1] = op2;

	switch (inst->type) {
	case INST_NOP: case INST_SAV:
	case INST_SWP: case INST_NEG:
		if (inst->opcnt != 0) return false;
		break;
	case INST_ADD: case INST_SUB:
		if (inst->opcnt != 1) return false;
		break;
	 case INST_JRO:
		if (inst->opcnt != 1) return false;
		if (inst->ops[0].type != OP_LIT) return false;
		break;
	case INST_JMP: case INST_JEZ: case INST_JNZ:
	case INST_JGZ: case INST_JLZ:
		if (inst->opcnt != 1) return false;
		if (inst->ops[0].type != OP_LABEL) return false;
		break;
	case INST_MOV:
		if (inst->opcnt != 2) return false;
		if (inst->ops[0].type == OP_LABEL) return false;
		if (inst->ops[1].type == OP_LABEL) return false;
		if (inst->ops[1].type == OP_LIT) return false;
		break;
	}

	return true;
}

struct tpu_inst *
tpu_add_inst(struct tpu *tpu, enum tpu_inst_type inst_type,
	unsigned opcnt, struct tpu_inst_op op1, struct tpu_inst_op op2)
{
	if (tpu->inst_cnt >= TPU_MAX_INST_CNT)
		die("tpu_add_inst: tpu X%i Y%i, >= max %i instructions",
			tpu->x, tpu->y, TPU_MAX_INST_CNT);
	if (!tpu_set_inst(tpu, (uint8_t) tpu->inst_cnt,
			inst_type, opcnt, op1, op2))
		return false;
	return &tpu->insts[tpu->inst_cnt++];
}

/* tpu can only read values if there is a writer attached and offering
 * a value -- this value is cleared on successful read */

static bool
tpu_port_read(struct tpu *tpu, enum tpu_port_dir dir, int *lit)
{
	struct tpu_port *port = &tpu->ports[dir];

	if (!port->avail)
		return false;

	*lit = port->in;
	port->reset_in = true;

	return true;
}

/* tpu can write value to ->out immediately but only finish the instruction
 * once the written value is read and cleared by the reader */

static bool
tpu_port_write(struct tpu *tpu, enum tpu_port_dir dir, int lit)
{
	struct tpu_port *port = &tpu->ports[dir];

	if (port->writing)
		return false;

	port->writing = true;
	port->out = lit;

	return true;
}

static void
tpu_jmp_label(struct tpu *tpu, const char *label)
{
	tpu->pc = label_map_get(&tpu->label_map, label);
	if (tpu->pc < 0 || tpu->pc >= tpu->inst_cnt) abort();
}

int
tpu_exec_get(struct tpu *tpu, struct tpu_inst_op *op, int *lit)
{
	int i;

	switch (op->type) {
	case OP_ACC:
		*lit = tpu->acc;
		return true;
	case OP_NIL:
		*lit = 0;
		return true;
	case OP_LEFT: case OP_RIGHT: case OP_UP: case OP_DOWN:
		return tpu_port_read(tpu, op_to_dir(op->type), lit);
	case OP_ANY:
		for (i = 0; i < 4; i++) {
			if (tpu_port_read(tpu, (enum tpu_port_dir) i, lit)) {
				tpu->last = i;
				return true;
			}
		}
		return false;
	case OP_LAST:
		if (tpu->last < 0) return false;
		return tpu_port_read(tpu, (enum tpu_port_dir) tpu->last, lit);
	case OP_LIT:
		*lit = op->val.lit;
		return true;
	case OP_LABEL:
		abort();
	}
}

bool
tpu_exec_put(struct tpu *tpu, struct tpu_inst_op *op, int lit)
{
	int i;

	switch (op->type) {
	case OP_ACC:
		tpu->acc = lit;
		return true;
	case OP_NIL:
		return true;
	case OP_LEFT: case OP_RIGHT: case OP_UP: case OP_DOWN:
		return tpu_port_write(tpu, op_to_dir(op->type), lit);
	case OP_ANY:
		for (i = 0; i < 4; i++) {
			if (tpu_port_write(tpu, (enum tpu_port_dir) i, lit)) {
				tpu->last = i;
				return true;
			}
		}
		return false;
	case OP_LAST:
		if (tpu->last < 0) return false;
		return tpu_port_write(tpu, (enum tpu_port_dir) tpu->last, lit);
	case OP_LABEL:
	case OP_LIT:
		abort();
	}
}

enum tpu_mode
tpu_exec(struct tpu *tpu, struct tpu_inst *inst)
{
	int lit;
	int val;

	tpu->idle = false;
	switch (inst->type) {
	case INST_NOP:
		tpu->pc += 1;
		return MODE_RUN;
	case INST_MOV:
		if (tpu->mode == MODE_WRITE) {
			tpu->idle = true;
			return MODE_WRITE;
		}
		if (!tpu_exec_get(tpu, &inst->ops[0], &lit)) {
			tpu->idle = true;
			return MODE_READ;
		}
		if (!tpu_exec_put(tpu, &inst->ops[1], lit)) {
			tpu->idle = true;
			return MODE_WRITE;
		}
		/* if MOV to port, set WRITE tentatively and fix on reset_in */
		if (inst->ops[1].type != OP_ACC && inst->ops[1].type != OP_NIL)
			return MODE_WRITE;
		tpu->pc += 1;
		return MODE_RUN;
	case INST_SWP:
		lit = tpu->acc;
		tpu->acc = tpu->bak;
		tpu->bak = lit;
		tpu->pc += 1;
		return MODE_RUN;
	case INST_SAV:
		tpu->bak = tpu->acc;
		tpu->pc += 1;
		return MODE_RUN;
	case INST_ADD:
		if (!tpu_exec_get(tpu, &inst->ops[0], &val)) {
			tpu->idle = true;
			return MODE_READ;
		}
		tpu->acc = MIN(MAX(tpu->acc + val, -999), 999);
		tpu->pc += 1;
		return MODE_RUN;
	case INST_SUB:
		 if (!tpu_exec_get(tpu, &inst->ops[0], &val)) {
			 tpu->idle = true;
			 return MODE_READ;
		 }
		tpu->acc = MIN(MAX(tpu->acc - val, -999), 999);
		tpu->pc += 1;
		return MODE_RUN;
	case INST_NEG:
		tpu->acc = -tpu->acc;
		tpu->pc += 1;
		return MODE_RUN;
	case INST_JMP:
		tpu_jmp_label(tpu, inst->ops[0].val.label);
		return MODE_RUN;
	case INST_JEZ:
		if (tpu->acc == 0)
			tpu_jmp_label(tpu, inst->ops[0].val.label);
		else
			tpu->pc += 1;
		return MODE_RUN;
	case INST_JNZ:
		if (tpu->acc != 0)
			tpu_jmp_label(tpu, inst->ops[0].val.label);
		else
			tpu->pc += 1;
		return MODE_RUN;
	case INST_JGZ:
		if (tpu->acc > 0)
			tpu_jmp_label(tpu, inst->ops[0].val.label);
		else
			tpu->pc += 1;
		return MODE_RUN;
	case INST_JLZ:
		if (tpu->acc < 0)
			tpu_jmp_label(tpu, inst->ops[0].val.label);
		else
			tpu->pc += 1;
		return MODE_RUN;
	case INST_JRO:
		tpu->pc += inst->ops[0].val.lit;
		if (tpu->pc < 0) tpu->pc = 0;
		if (tpu->pc >= tpu->inst_cnt)
			tpu->pc = (int) tpu->inst_cnt - 1;
		return MODE_RUN;
	default:
		abort();
	}
}

void
tpu_step(struct tpu *tpu)
{
	int prev_mode;

	prev_mode = tpu->mode;
	if (tpu->pc < tpu->inst_cnt) {
		tpu->mode = tpu_exec(tpu, &tpu->insts[tpu->pc]);
		if (tpu->pc >= tpu->inst_cnt) tpu->pc = 0;
	} else {
		tpu->mode = MODE_IDLE;
		tpu->idle = true;
	}
}

void
tpu_map_init(struct tpu_map *map)
{
	memset(map->buckets, 0, sizeof(void *) * TPU_MAP_BUCKETS);
}

void
tpu_map_deinit(struct tpu_map *map)
{
	struct tpu_map_link *link, *next;
	size_t i;

	for (i = 0; i < TPU_MAP_BUCKETS; i++) {
		link = map->buckets[i];
		while (link) {
			next = link->next;
			tpu_deinit(link->tpu);
			free(link->tpu);
			free(link);
			link = next;
		}
	}
}

static struct tpu_map_link **
tpu_map_link_pos(struct tpu_map *map, int x, int y)
{
	struct tpu_map_link **link;
	size_t i;

	i = (size_t) (x + y) % TPU_MAP_BUCKETS;
	link = &map->buckets[i];
	while (*link && !((*link)->x == x && (*link)->y == y))
		link = &(*link)->next;

	return link;
}

bool
tpu_map_add(struct tpu_map *map, struct tpu *tpu)
{
	struct tpu_map_link **pos, *link;

	pos = tpu_map_link_pos(map, tpu->x, tpu->y);
	if (*pos) return false;
	*pos = link = malloc(sizeof(struct tpu_map_link));
	if (!link) die("malloc:");
	link->tpu = tpu;
	link->x = tpu->x;
	link->y = tpu->y;
	link->next = NULL;
	return true;
}

struct tpu *
tpu_map_get(struct tpu_map *map, int x, int y)
{
	struct tpu_map_link **link;

	if (x < 0 || y < 0) return NULL;

	link = tpu_map_link_pos(map, x, y);
	if (!*link) return NULL;

	return (*link)->tpu;
}

void
tis_init(struct tis *tis)
{
	tpu_map_init(&tis->tpu_map);
	memset(&tis->in_ports, 0, TIS_MAX_IO_PORTS * sizeof(void *));
	memset(&tis->out_ports, 0, TIS_MAX_IO_PORTS * sizeof(void *));
	tis->steps = 0;
	tis->show_stats = false;
	tis->asm_filepath = NULL;
}

void
tis_deinit(struct tis *tis)
{
	int i;

	tpu_map_deinit(&tis->tpu_map);
	for (i = 0; i < TIS_MAX_IO_PORTS; i++) {
		if (tis->in_ports[i]) {
			tpu_io_port_deinit(tis->in_ports[i]);
			free(tis->in_ports[i]);
		}
		if (tis->out_ports[i]) {
			tpu_io_port_deinit(tis->out_ports[i]);
			free(tis->out_ports[i]);
		}
	}
	free(tis->asm_filepath);
}

bool
tis_step(struct tis *tis)
{
	struct tpu_map_link *link;
	bool running;
	size_t i;

	for (i = 0; i < TPU_MAP_BUCKETS; i++) {
		link = tis->tpu_map.buckets[i];
		for (; link; link = link->next)
			tpu_step(link->tpu);
	}

	tis_communicate(tis);

	for (i = 0; i < TPU_MAP_BUCKETS; i++) {
		link = tis->tpu_map.buckets[i];
		for (; link; link = link->next)
			tpu_update(link->tpu);
	}

	for (i = 0; i < TIS_MAX_IO_PORTS; i++) {
		if (tis->in_ports[i])
			tpu_port_update(&tis->in_ports[i]->port);
		if (tis->out_ports[i])
			tpu_port_update(&tis->out_ports[i]->port);
	}

	running = false;
	for (i = 0; i < TPU_MAP_BUCKETS; i++) {
		link = tis->tpu_map.buckets[i];
		for (; link; link = link->next) {
			running |= !link->tpu->idle;
			if (link->tpu->pc < link->tpu->inst_cnt) {
				link->tpu->steps += 1;
				if (link->tpu->idle)
					link->tpu->idle_steps += 1;
			}
		}
	}

	tis->steps += 1;

	return running;
}

void
tis_communicate(struct tis *tis)
{
	struct tpu_io_port *io_port;
	char buf[16], *s;
	int val, i;

	for (i = 0; i < TIS_MAX_IO_PORTS; i++) {
		io_port = tis->in_ports[i];
		if (!io_port) continue;
		if (!io_port->port.attached || io_port->port.writing)
			continue;
		if (!io_port->file) continue;
		while ((s = fgets(buf, sizeof(buf), io_port->file))) {
			if ((s = strchr(buf, '\n'))) *s = '\0';
			if (!*buf) continue;
			val = (int) strtol(buf, &s, 10);
			if (!s || *s != '\0')
				die("communicate: invalid input '%s'", buf);
			io_port->port.out = val;
			io_port->port.writing = true;
			io_port->io_step = tis->steps;
			break;
		}
	}

	for (i = 0; i < TIS_MAX_IO_PORTS; i++) {
		io_port = tis->out_ports[i];
		if (!io_port) continue;
		if (!io_port->port.attached || !io_port->port.avail)
			continue;
		if (io_port->file)
			fprintf(io_port->file, "%i\n", io_port->port.in);
		io_port->port.reset_in = true;
		tpu_port_update(&io_port->port);
		io_port->io_step = tis->steps;
	}
}

struct tis_stats
tis_gen_stats(struct tis *tis)
{
	struct tis_stats stats;
	struct tpu_map_link *link;
	size_t idle_steps;
	size_t all_steps;
	int i;

	stats.steps = 0;
	stats.nodes = 0;
	stats.insts = 0;
	stats.idle = 0;

	for (i = 0; i < TIS_MAX_IO_PORTS; i++) {
		if (tis->out_ports[i] && tis->out_ports[i]->io_step > 0) {
			stats.steps = MAX(stats.steps, tis->out_ports[i]->io_step + 1);
		}
	}

	all_steps = 0;
	idle_steps = 0;
	for (i = 0; i < TPU_MAP_BUCKETS; i++) {
		link = tis->tpu_map.buckets[i];
		for (; link; link = link->next) {
			stats.insts += link->tpu->inst_cnt;
			idle_steps += link->tpu->idle_steps;
			all_steps += link->tpu->steps;
			stats.nodes += (link->tpu->inst_cnt > 0);
		}
	}

	if (!stats.steps) stats.steps = all_steps;

	if (all_steps == 0)
		stats.idle = 0;
	else
		stats.idle = (float) idle_steps / (float) all_steps;

	return stats;
}