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

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

enum {
	MOVE_NORTH = 1,
	MOVE_SOUTH,
	MOVE_WEST,
	MOVE_EAST
};

enum {
	BLOCK_WALL = 0,
	BLOCK_OPEN,
	BLOCK_GOAL
};

struct state {
	int back;
	int next;
};

const char block_lut[] = {
	'#',
	'.',
	'O'
};

const struct vec2i dirs[] = {
	{ 0 },
	{ 0, -1 },
	{ 0, 1 },
	{ -1, 0 },
	{ 1, 0 },
};

const int rev[] = {
	0, 2, 1, 4, 3
};

uint32_t
vec_hmap_hash(struct hmap_key key)
{
	const struct vec2i *vec = key.p;

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

bool
vec_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;
}

int
robot_move(struct icc *icc, int dir)
{
	while (icc->state != ICC_HALT) {
		icc_step_inst(icc);
		switch (icc->state) {
		case ICC_INPUT:
			mi_setv(&icc->in, (mi_ul) dir, MI_POS);
			continue;
		case ICC_OUTPUT:
			return (int) mi_cast_ul(&icc->out);
		case ICC_HALT:
		case ICC_RUN:
			continue;
		}
	}

	return -1;
}

void
print_map(struct hmap *map, struct vec2i droid)
{
	struct hmap_link *link;
	struct hmap_iter iter;
	struct vec2i min = { 0 };
	struct vec2i max = { 0 };
	struct vec2i pos;
	const struct vec2i *vec;

	for (HMAP_ITER(map, iter)) {
		vec = iter.link->key.p;
		min.x = MIN(min.x, vec->x);
		min.y = MIN(min.y, vec->y);
		max.x = MAX(max.x, vec->x);
		max.y = MAX(max.y, vec->y);
	}

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

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

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

size_t
dijkstra_dist(struct hmap *map, struct vec2i start, const struct vec2i *end)
{
	struct hmap_link *link, **linkp;
	struct hmap_iter iter;
	struct hmap seen;
	struct dvec outer;
	struct vec2i *vec;
	struct vec2i pos, npos;
	size_t dist, maxdist;
	void *key;
	int dir;

	hmap_init(&seen, 32, vec_hmap_hash, vec_hmap_keycmp,
		&stdlib_strict_heap_allocator);
	dvec_init(&outer, sizeof(struct vec2i),
		10, &stdlib_strict_heap_allocator);

	vec = dvec_add_slot(&outer);
	vec2i_set(vec, &start);

	key = memdup(vec, sizeof(struct vec2i));
	hmap_add(&seen, (struct hmap_key) {.p = key},
		(struct hmap_val) {.s = 0});

	maxdist = 0;
	while (!dvec_empty(&outer)) {
		vec = dvec_back(&outer);
		vec2i_set(&pos, vec);
		link = hmap_get(&seen, (struct hmap_key) {.p = &pos});
		assert(link);
		dist = link->value.s + 1;
		dvec_rm_slot(&outer, vec);

		for (dir = 1; dir <= 4; dir++) {
			vec2i_add(&npos, &pos, &dirs[dir]);
			linkp = hmap_link_pos(&seen,
				(struct hmap_key) {.p = &npos});
			if (*linkp) continue;

			if (end && !memcmp(&npos, end, sizeof(struct vec2i)))
				goto exit;

			link = hmap_get(map, (struct hmap_key) {.p = &npos});
			assert(link);
			if (link->value.c != '.') continue;

			if (!maxdist || dist > maxdist)
				maxdist = dist;

			vec = dvec_add_slot(&outer);
			vec2i_set(vec, &npos);
			key = memdup(&npos, sizeof(struct vec2i));
			*linkp = hmap_link_alloc(&seen,
				(struct hmap_key) {.p = key},
				(struct hmap_val) {.s = dist}, NULL);
		}
	}

exit:
	for (HMAP_ITER(&seen, iter))
		free(iter.link->key._p);
	dvec_deinit(&outer);
	hmap_deinit(&seen);

	return end ? dist : maxdist;
}

void
explore_map(struct icc *icc, struct hmap *map,
	struct vec2i start, struct vec2i *tank)
{
	struct hmap_link **linkp;
	struct hmap_link *link;
	struct dvec states;
	struct vec2i pos, npos;
	struct state *state;
	int move_result;
	int move_dir;
	void *key;

	dvec_init(&states, sizeof(struct state),
		10, &stdlib_strict_heap_allocator);

	state = dvec_add_slot(&states);
	state->back = -1;
	state->next = 1;

	vec2i_set(&pos, &start);
	key = memdup(&pos, sizeof(struct vec2i));
	hmap_add(map, (struct hmap_key) {.p = key},
		(struct hmap_val) {.c = '.'});

	while (!dvec_empty(&states)) {
		assert(!dvec_empty(&states));
		state = dvec_back(&states);
		for (; state->next <= 4; state->next++) {
			if (state->next == state->back)
				continue;
			move_dir = state->next;
			vec2i_add(&npos, &pos, &dirs[move_dir]);
			link = hmap_get(map,
				(struct hmap_key) {.p = &npos});
			if (!link) break;
		}

		if (state->next > 4) {
			if (state->back <= 0) break;
			move_result = robot_move(icc, state->back);
			assert(move_result == BLOCK_OPEN);
			vec2i_add(&pos, &pos, &dirs[state->back]);
			dvec_rm_slot(&states, state);
			if (aoc.debug) {
				aoc_debug("Step: backtrack x %i y %i\n", pos.x, pos.y);
				print_map(map, pos);
				aoc_debug("\n");
			}
			continue;
		}
		state->next += 1;

		assert(move_dir >= 1 && move_dir <= 4);
		move_result = robot_move(icc, move_dir);
		assert(move_result >= 0 && move_result <= 2);

		linkp = hmap_link_pos(map, (struct hmap_key) {.p = &npos});
		if (*linkp) {
			(*linkp)->value.c = block_lut[move_result];
		} else {
			key = memdup(&npos, sizeof(struct vec2i));
			*linkp = hmap_link_alloc(map,
				(struct hmap_key) {.p = key},
				(struct hmap_val) {.c = block_lut[move_result]},
				NULL);
		}

		switch (move_result) {
		case BLOCK_GOAL:
			vec2i_set(tank, &npos);
		case BLOCK_OPEN:
			vec2i_set(&pos, &npos);
			state = dvec_add_slot(&states);
			state->back = rev[move_dir];
			state->next = 1;
			break;
		case BLOCK_WALL:
			break;
		default:
			assert(0);
		}

		if (aoc.debug) {
			aoc_debug("Step: x %i y %i dx %i dy %i\n",
				pos.x, pos.y, dirs[move_dir].x, dirs[move_dir].y);
			print_map(map, pos);
			aoc_debug("\n");
		}
	}

	dvec_deinit(&states);
}

void
part1(void)
{
	struct icc icc;
	struct hmap_iter iter;
	struct hmap map;
	struct vec2i start, tank;
	size_t dist;

	icc_init(&icc);
	hmap_init(&map, 32, vec_hmap_hash, vec_hmap_keycmp,
		&stdlib_strict_heap_allocator);

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

	start = (struct vec2i) { 0, 0 };
	explore_map(&icc, &map, start, &tank);

	dist = dijkstra_dist(&map, start, &tank);

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

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

void
part2(void)
{
	struct icc icc;
	struct hmap_iter iter;
	struct hmap map;
	struct vec2i start, tank;
	size_t dist;

	icc_init(&icc);
	hmap_init(&map, 32, vec_hmap_hash, vec_hmap_keycmp,
		&stdlib_strict_heap_allocator);

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

	start = (struct vec2i) { 0, 0 };
	explore_map(&icc, &map, start, &tank);

	dist = dijkstra_dist(&map, tank, NULL);

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

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