#include "aoc.h"
#include "allocator.h"
#include "util.h"
#include "iccmp.h"
#include "hmap.h"
#include "dvec_s.h"
#include "vec.h"

#include <assert.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>

enum {
	DIR_NORTH,
	DIR_EAST,
	DIR_SOUTH,
	DIR_WEST,
};

struct path {
	struct vec2i start, end;
};

static const struct vec2i dirs[] = {
	[DIR_NORTH] = {  0, -1 },
	[DIR_EAST]  = {  1,  0 },
	[DIR_SOUTH] = {  0,  1 },
	[DIR_WEST]  = { -1,  0 },
};

bool
vec2i_hmap_keycmp(struct hmap_key k1, struct hmap_key k2)
{
	const struct vec2i *v1 = k1.p, *v2 = k2.p;

	return v1->x == v2->x && v1->y == v2->y;
}

uint32_t
vec2i_hmap_hash(struct hmap_key key)
{
	const struct vec2i *v = key.p;

	return (uint32_t) (v->x << 16) ^ (uint32_t) v->y;
}

int
dir_from_char(char c)
{
	switch (c) {
	case '^':
		return DIR_NORTH;
	case '>':
		return DIR_EAST;
	case 'v':
		return DIR_SOUTH;
	case '<':
		return DIR_WEST;
	default:
		return -1;
	}
}

int
dir_from_path(struct path *path)
{
	if (path->end.x > path->start.x)
		return DIR_EAST;
	else if (path->end.x < path->start.x)
		return DIR_WEST;
	else if (path->end.y > path->start.y)
		return DIR_SOUTH;
	else if (path->end.y < path->start.y)
		return DIR_NORTH;
	assert(0);
}

void
receive_map(struct icc *icc, struct hmap *map,
	struct vec2i *min, struct vec2i *max,
	struct vec2i *dpos, int *ddir)
{
	struct hmap_iter iter;
	struct vec2i pos;
	void *key;
	char c;

	for (HMAP_ITER(map, iter))
		free(iter.link->key._p);
	hmap_clear(map);

	vec2i_setv(min, 0, 0);
	vec2i_setv(max, 0, 0);
	vec2i_setv(&pos, 0, 0);
	while (icc->state != ICC_HALT) {
		icc_step_inst(icc);
		if (icc->state == ICC_INPUT)
			break;
		if (icc->state == ICC_OUTPUT) {
			c = (char) mi_cast_ul(&icc->out);
			if (c == '\n') {
				if (pos.x == 0) break;
				pos.y += 1;
				pos.x = 0;
			} else {
				if (dir_from_char(c) >= 0) {
					vec2i_set(dpos, &pos);
					*ddir = dir_from_char(c);
				}
				key = memdup(&pos, sizeof(struct vec2i));
				hmap_add(map, (struct hmap_key) {.p = key},
					(struct hmap_val) {.c = c});
				if (pos.x > max->x) max->x = pos.x;
				if (pos.y > max->y) max->y = pos.y;
				if (pos.x < min->x) min->x = pos.x;
				if (pos.y < min->y) min->y = pos.y;
				pos.x += 1;
			}
		}
	}
}

void
receive_output(struct icc *icc)
{
	aoc_debug("ICC> ");
	while (icc->state != ICC_HALT) {
		icc_step_inst(icc);
		if (icc->state == ICC_INPUT)
			break;
		if (icc->state == ICC_OUTPUT)
			aoc_debug("%c", (char) mi_cast_ul(&icc->out));
	}
	assert(icc->state != ICC_HALT);
}

void
print_map(struct hmap *map, struct vec2i min, struct vec2i max)
{
	struct hmap_link *link;
	struct vec2i pos;

	aoc_debug("+");
	for (pos.x = min.x; pos.x <= max.x; pos.x++)
		aoc_debug("-");
	aoc_debug("+\n");

	for (pos.y = min.y; pos.y <= max.y; pos.y++) {
		aoc_debug("|");
		for (pos.x = min.x; pos.x <= max.x; pos.x++) {
			link = hmap_get(map, (struct hmap_key) {.p = &pos});
			if (link) aoc_debug("%c", link->value.c);
			else aoc_debug(" ");
		}
		aoc_debug("|\n");
	}

	aoc_debug("+");
	for (pos.x = min.x; pos.x <= max.x; pos.x++)
		aoc_debug("-");
	aoc_debug("+\n");
}

void
find_path_end(struct hmap *map, struct vec2i start,
	size_t dir, struct vec2i *end)
{
	struct hmap_link *link;
	struct vec2i ppos, pos;

	vec2i_set(&pos, &start);
	while (1) {
		vec2i_set(&ppos, &pos);
		vec2i_add(&pos, &pos, &dirs[dir]);
		link = hmap_get(map, (struct hmap_key) {.p = &pos});
		if (!link || link->value.c != '#')
			break;
	}

	vec2i_set(end, &ppos);
}

void
add_path(struct hmap *paths, struct path *path)
{
	struct hmap_link **link;
	struct dvec *list;
	void *key;

	link = hmap_link_pos(paths, (struct hmap_key) {.p = &path->start});
	if (*link) {
		list = (*link)->value._p;
		memcpy(dvec_add_slot(list), path, sizeof(struct path));
	} else {
		key = memdup(&path->start, sizeof(struct vec2i));
		list = dvec_alloc(sizeof(struct path),
			1, &stdlib_strict_heap_allocator, NULL);
		memcpy(dvec_add_slot(list), path, sizeof(struct path));
		hmap_add(paths,
			(struct hmap_key) {.p = key},
			(struct hmap_val) {.p = list});
	}
}

void
gen_paths(struct hmap *paths, struct hmap *map,
	struct vec2i min, struct vec2i max)
{
	struct hmap_link *link;
	struct vec2i pos, npos;
	struct vec2i end;
	struct path path;
	bool side[4];
	bool endpoint;
	size_t i;

	vec2i_setv(&pos, 0, 0);
	for (pos.y = min.y; pos.y <= max.y; pos.y++) {
		for (pos.x = min.x; pos.x <= max.x; pos.x++) {
			link = hmap_get(map, (struct hmap_key) {.p = &pos});
			if (!link || link->value.c == '.')
				continue;

			for (i = 0; i < ARRLEN(dirs); i++) {
				side[i] = false;
				vec2i_add(&npos, &pos, &dirs[i]);
				link = hmap_get(map, (struct hmap_key) {.p = &npos});
				side[i] = link ? link->value.c == '#' : false;
			}

			endpoint = (side[DIR_NORTH] != side[DIR_SOUTH]);
			endpoint |= (side[DIR_WEST] != side[DIR_EAST]);
			if (endpoint) {
				aoc_debug("node %li %li\n", pos.x, pos.y);
				for (i = 0; i < ARRLEN(dirs); i++) {
					if (!side[i]) continue;
					find_path_end(map, pos, i, &end);
					vec2i_set(&path.start, &pos);
					vec2i_set(&path.end, &end);
					add_path(paths, &path);
				}
			}
		}
	}
}

void
gen_insts(struct dvec *insts, struct hmap *paths,
	struct vec2i start, int ddir)
{
	struct hmap_link *link;
	struct dvec *list;
	struct path *path;
	struct vec2i ppos, pos;
	char *inst, type;
	int cnt, dir, ndir;
	bool found;

	dir = ddir;
	vec2i_set(&pos, &start);
	vec2i_set(&ppos, &pos);
	while (1) {
		link = hmap_get(paths, (struct hmap_key) {.p = &pos});
		assert(link != NULL);

		found = false;
		list = link->value._p;
		for (DVEC_ITER(list, path)) {
			if (vec2i_eql(&ppos, &path->end))
				continue;
			found = true;
			break;
		}

		if (!found) break;

		aoc_debug("path %i %i -> %i %i\n", path->start.x, path->start.y,
			path->end.x, path->end.y);

		vec2i_set(&ppos, &pos);
		vec2i_set(&pos, &path->end);

		ndir = dir_from_path(path);
		aoc_debug("ndir %i\n", ndir);
		type = ndir > dir ? 'R' : 'L';
		cnt = ABS(ndir - dir);
		if (cnt > 2) {
			cnt = 4 - cnt;
			type = type == 'L' ? 'R' : 'L';
		}
		for (; cnt > 0; cnt--) {
			inst = dvec_add_slot(insts);
			*inst = type;
		}
		dir = ndir;

		cnt = (int) ABS(path->end.x - path->start.x)
			+ (int) ABS(path->end.y - path->start.y);
		inst = dvec_add_slot(insts);
		*inst = (char) -cnt;
	}
}

char *
insts_line(struct dvec *insts)
{
	char *str;
	char *c;

	str = NULL;
	for (DVEC_ITER(insts, c)) {
		if (str) str = apprintf(str, ",");
		if (*c < 0) {
			str = apprintf(str, "%i", -*c);
		} else {
			str = apprintf(str, "%c", *c);
		}
	}
	str = apprintf(str, "\n");

	return str;
}

size_t
gen_func(struct dvec *insts, struct dvec *main,
	struct dvec *cur, char cc, size_t start, size_t len,
	struct dvec *f1, char c1, struct dvec *f2, char c2)
{
	char *inst;
	size_t i;

	for (i = start; i < MIN(insts->len, start + len); i++) {
		inst = dvec_add_slot(cur);
		*inst = *(char *)dvec_at(insts, i);
	}

	inst = dvec_add_slot(main);
	*inst = cc;

	while (1) {
		if (f1 && i + f1->len <= insts->len
				&& !memcmp(dvec_at(insts, i), f1->data, f1->len)) {
			i += f1->len;
			inst = dvec_add_slot(main);
			*inst = c1;
		} else if (f2 && i + f2->len <= insts->len
				&& !memcmp(dvec_at(insts, i), f2->data, f2->len)) {
			i += f2->len;
			inst = dvec_add_slot(main);
			*inst = c2;
		} else if (i + cur->len <= insts->len
				&& !memcmp(dvec_at(insts, i), cur->data, cur->len)) {
			i += cur->len;
			inst = dvec_add_slot(main);
			*inst = cc;
		} else {
			break;
		}
	}

	return i;
}

void
gen_funcs(struct dvec *insts, struct dvec *main,
	struct dvec *fa, struct dvec *fb, struct dvec *fc)
{
	size_t alen, blen, clen;
	size_t lens, pos;
	char *line;

	lens = 0;
	while (lens < 20 * 20 * 20) {
		dvec_clear(main);
		dvec_clear(fa);
		dvec_clear(fb);
		dvec_clear(fc);

		lens += 1;
		alen = (lens / 20 / 20) % 20;
		blen = (lens / 20) % 20;
		clen = lens % 20;
		if (!alen || !blen || !clen)
			continue;

		pos = 0;
		pos = gen_func(insts, main, fa, 'A', pos, alen, NULL, 0, NULL, 0);
		pos = gen_func(insts, main, fb, 'B', pos, blen, fa, 'A', NULL, 0);
		pos = gen_func(insts, main, fc, 'C', pos, clen, fa, 'A', fb, 'B');

		assert(fa->len == alen);
		assert(fb->len == blen);
		assert(fc->len == clen);

		if (aoc.debug && main->len > 3) {
			line = insts_line(main);
			aoc_debug("%2li %2li | %2li %2li %2li %s",
				pos, insts->len, alen, blen, clen, line);
			free(line);

			line = insts_line(insts);
			aoc_debug("%s", line);
			free(line);

			line = insts_line(fa);
			aoc_debug("%s", line);
			free(line);

			line = insts_line(fb);
			aoc_debug("%s", line);
			free(line);

			line = insts_line(fc);
			aoc_debug("%s", line);
			free(line);

			aoc_debug("\n");
		}

		if (pos == insts->len && main->len <= 20)
			return;
	}
	assert(0);
}

void
send_line(struct icc *icc, const char *line)
{
	const char *c;

	c = line;
	while (icc->state != ICC_HALT && *c) {
		icc_step_inst(icc);
		switch (icc->state) {
		case ICC_HALT:
		case ICC_RUN:
			continue;
		case ICC_INPUT:
			mi_setv(&icc->in, (mi_ul) *c, MI_POS);
			c += 1;
			break;
		default:
			receive_output(icc);
			assert(0);
		}
	}
	assert(icc->state != ICC_HALT);
}

void
part1(void)
{
	struct icc icc;
	struct hmap map;
	struct hmap_link *link;
	struct hmap_iter iter;
	struct vec2i pos, npos;
	struct vec2i min, max;
	struct vec2i dpos;
	size_t answer;
	size_t i;
	int ddir;

	icc_init(&icc);
	hmap_init(&map, 32, vec2i_hmap_hash, vec2i_hmap_keycmp,
		&stdlib_strict_heap_allocator);

	icc_parse_inst(&icc, aoc.input, aoc.input_size);

	receive_map(&icc, &map, &min, &max, &dpos, &ddir);

	if (aoc.debug)
		print_map(&map, min, max);

	vec2i_setv(&pos, 0, 0);
	answer = 0;
	for (pos.y = min.y; pos.y <= max.y; pos.y++) {
		for (pos.x = min.x; pos.x <= max.x; pos.x++) {
			link = hmap_get(&map, (struct hmap_key) {.p = &pos});
			if (!link || link->value.c == '.')
				continue;
			for (i = 0; i < ARRLEN(dirs); i++) {
				vec2i_add(&npos, &pos, &dirs[i]);
				link = hmap_get(&map, (struct hmap_key) {.p = &npos});
				if (!link || link->value.c != '#')
					break;
			}
			if (i == ARRLEN(dirs))
				answer += (size_t) pos.x * (size_t) pos.y;
		}
	}

	aoc.answer = aprintf("%lu", answer);
	aoc.solution = "4112";

	for (HMAP_ITER(&map, iter))
		free(iter.link->key._p);
	hmap_deinit(&map);
	icc_deinit(&icc);
}

void
part2(void)
{
	struct icc icc;
	struct hmap map;
	struct hmap paths;
	struct dvec insts;
	struct dvec main;
	struct dvec fa, fb, fc;
	struct hmap_iter iter;
	struct vec2i min, max;
	struct vec2i dpos;
	struct maxint mi_addr = MI_INIT;
	struct maxint mi_val = MI_INIT;
	mi_ul ul_addr, ul_val;
	char *line, buf[256];
	int ddir;

	icc_init(&icc);
	hmap_init(&map, 32, vec2i_hmap_hash, vec2i_hmap_keycmp,
		&stdlib_strict_heap_allocator);
	hmap_init(&paths, 32, vec2i_hmap_hash, vec2i_hmap_keycmp,
		&stdlib_strict_heap_allocator);
	dvec_init(&insts, sizeof(char), 32, &stdlib_strict_heap_allocator);
	dvec_init(&main, sizeof(char), 0, &stdlib_strict_heap_allocator);
	dvec_init(&fa, sizeof(char), 0, &stdlib_strict_heap_allocator);
	dvec_init(&fb, sizeof(char), 0, &stdlib_strict_heap_allocator);
	dvec_init(&fc, sizeof(char), 0, &stdlib_strict_heap_allocator);

	icc_parse_inst(&icc, aoc.input, aoc.input_size);

	icc_write(&icc, mi_imm(&mi_addr, &ul_addr, 0, MI_POS),
		mi_imm(&mi_val, &ul_val, 2, MI_POS));

	receive_map(&icc, &map, &min, &max, &dpos, &ddir);

	if (aoc.debug)
		print_map(&map, min, max);

	gen_paths(&paths, &map, min, max);

	gen_insts(&insts, &paths, dpos, ddir);

	if (aoc.debug) {
		aoc_debug("%li insts\n", insts.len);
		line = insts_line(&insts);
		aoc_debug("%s", line);
		free(line);
	}

	gen_funcs(&insts, &main, &fa, &fb, &fc);

	receive_output(&icc);
	line = insts_line(&main);
	if (aoc.debug) printf("main %s", line);
	send_line(&icc, line);
	free(line);

	receive_output(&icc);
	line = insts_line(&fa);
	if (aoc.debug) printf("A %s", line);
	send_line(&icc, line);
	free(line);

	receive_output(&icc);
	line = insts_line(&fb);
	if (aoc.debug) printf("B %s", line);
	send_line(&icc, line);
	free(line);

	receive_output(&icc);
	line = insts_line(&fc);
	if (aoc.debug) printf("C %s", line);
	send_line(&icc, line);
	free(line);

	receive_output(&icc);
	send_line(&icc, "n\n");

	while (icc.state != ICC_HALT)
		icc_step_inst(&icc);
exit:
	mi_dec(&icc.out, buf, sizeof(buf), 0);
	aoc.answer = aprintf("%s", buf);
	aoc.solution = "578918";

	dvec_deinit(&fc);
	dvec_deinit(&fb);
	dvec_deinit(&fa);
	dvec_deinit(&main);
	dvec_deinit(&insts);

	for (HMAP_ITER(&paths, iter)) {
		free(iter.link->key._p);
		dvec_free(iter.link->value._p);
	}
	hmap_deinit(&paths);

	for (HMAP_ITER(&map, iter))
		free(iter.link->key._p);
	hmap_deinit(&map);

	icc_deinit(&icc);
}