#include "allocator.h"
#include "aoc.h"
#include "hmap.h"
#include "hmap_s.h"
#include "vec_s.h"
#include "util.h"
#include "list.h"

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

#define PORTAL_ID(a, b, i) \
	((size_t) (((uint8_t) (i) << 16) + ((uint8_t) (a) << 8) + (uint8_t) (b)))
#define PORTAL_CC_MASK 0xffff
#define PORTAL_IN_MASK 0x10000

struct djk_state {
	size_t dist;
	struct vec3i pos;

	struct list_link link;
};

void
vec3i_add_2d(struct vec3i *dst, const struct vec3i *src, const struct vec2i *off)
{
	dst->x = src->x + off->x;
	dst->y = src->y + off->y;
	dst->z = src->z;
}

static bool
djk_dist_ordering(const void *p1, const void *p2, void *u)
{
	const struct djk_state *s1 = p1, *s2 = p2;

	return s1->dist < s2->dist;
}

static void
load_map(struct hmap *map, const char *input, size_t len,
	struct vec2i *min, struct vec2i *max)
{
	char line[256];
	const char *lpos, *lend;
	struct vec2i pos, *key;

	lpos = aoc.input;
	lend = lpos + aoc.input_size;
	max->x = max->y = 0;
	vec2i_setv(&pos, 0, 0);
	while (readtok(line, sizeof(line), '\n', &lpos, lend)) {
		for (pos.x = 0; line[pos.x]; pos.x++) {
			if (line[pos.x] == ' ')
				continue;
			if (line[pos.x] == '.') {
				min->x = min->x ? MIN(min->x, pos.x) : pos.x;
				min->y = min->y ? MIN(min->y, pos.y) : pos.y;
				max->x = max->x ? MAX(max->x, pos.x) : pos.x;
				max->y = max->y ? MAX(max->y, pos.y) : pos.y;
			}
			key = malloc(sizeof(struct vec3i));
			vec2i_set(key, &pos);
			hmap_add(map, (struct hmap_key) {.p = key},
				(struct hmap_val) {.c = line[pos.x]});
		}
		pos.y += 1;
	}
}

static void
load_portals(struct hmap *portals, struct hmap *map,
	struct vec2i *min, struct vec2i *max)
{
	struct hmap_link *link1;
	struct hmap_link *link2;
	struct hmap_iter iter;
	struct vec2i pos;
	struct vec2i *key;
	size_t i, id;
	bool inner;

	for (HMAP_ITER(map, iter)) {
		if (!isupper(iter.link->value.c))
			continue;

		for (i = 0; i < 4; i++) {
			vec2i_add(&pos, iter.link->key.p, &adj[i]);
			link1 = hmap_get(map, (struct hmap_key) {.p = &pos});
			if (link1 && isupper(link1->value.c))
				break;
		}
		if (i != ADJ_SOUTH && i != ADJ_EAST)
			continue;

		for (i = 0; i < 4; i++) {
			vec2i_add(&pos, iter.link->key.p, &adj[i]);
			link2 = hmap_get(map, (struct hmap_key) {.p = &pos});
			if (link2 && link2->value.c == '.')
				break;
		}
		if (i == 4) {
			for (i = 0; i < 4; i++) {
				vec2i_add(&pos, link1->key.p, &adj[i]);
				link2 = hmap_get(map, (struct hmap_key) {.p = &pos});
				if (link2 && link2->value.c == '.')
					break;
			}
		}
		assert(i != 4);
		vec2i_set(&pos, link2->key.p);

		inner = pos.x != min->x && pos.y != min->y
			&& pos.x != max->x && pos.y != max->y;
		aoc_debug("PORTAL %c %c %li %li %i\n",
			iter.link->value.c, link1->value.c,
			pos.x, pos.y, inner);

		key = memdup(&pos, sizeof(struct vec2i));
		id = PORTAL_ID(iter.link->value.c, link1->value.c, inner);
		hmap_add(portals, (struct hmap_key) {.p = key},
			(struct hmap_val) {.s = id});
	}
}

static bool
get_portal(struct hmap *portals, struct vec3i *pos, size_t id)
{
	struct hmap_iter iter;
	const struct vec3i *tmp;

	for (HMAP_ITER(portals, iter)) {
		tmp = iter.link->key.p;
		if (iter.link->value.s == id) {
			pos->x = tmp->x;
			pos->y = tmp->y;
			pos->z = 0;
			return true;
		}
	}

	return false;
}

static void
add_djk_state(struct list *active, struct hmap *distmap,
	struct hmap *portals, struct vec3i *pos, size_t dist)
{
	struct hmap_link **hmap_linkp;
	struct djk_state *nstate;
	struct vec3i *key;
	bool new;

	hmap_linkp = hmap_link_pos(distmap, (struct hmap_key) {.p = pos});
	if (!*hmap_linkp) {
		key = memdup(pos, sizeof(struct vec3i));
		*hmap_linkp = hmap_link_alloc(distmap,
			(struct hmap_key) {.p = key},
			(struct hmap_val) {.s = dist}, NULL);
		new = true;
	} else {
		if (dist < (*hmap_linkp)->value.s) {
			(*hmap_linkp)->value.s = dist;
			new = true;
		} else {
			new = false;
		}
	}

	if (new) {
		aoc_debug("ADDING NEW\n");
		nstate = malloc(sizeof(struct djk_state));
		nstate->dist = dist;
		vec3i_set(&nstate->pos, pos);
		nstate->link = LIST_LINK_INIT;
		list_insert_sorted(active, &nstate->link,
			false, djk_dist_ordering,
			LIST_OFFSET(struct djk_state, link), NULL);
	}
}

static size_t
min_dist(struct hmap *map, struct hmap *portals,
	struct vec3i *start, struct vec3i *end)
{
	struct list active;
	struct hmap distmap;
	struct list_link *list_link;
	struct hmap_link *hmap_link;
	struct hmap_iter hmap_iter;
	struct djk_state *state;
	struct vec3i *key, npos;
	size_t mdist, i, id;

	list_init(&active);
	hmap_init(&distmap, 256, vec3i_hmap_hash, vec3i_hmap_keycmp,
		&stdlib_strict_heap_allocator);

	state = malloc(sizeof(struct djk_state));
	state->dist = 0;
	vec3i_set(&state->pos, start);
	list_insert_front(&active, &state->link);

	key = memdup(start, sizeof(struct vec3i));
	hmap_add(&distmap, (struct hmap_key) {.p = key},
		(struct hmap_val) {.s = 0});

	mdist = 0;
	while (!list_empty(&active)) {
		list_link = list_front(&active);
		state = LIST_UPCAST(list_link, struct djk_state, link);
		if (vec3i_eql(&state->pos, end)) {
			mdist = state->dist;
			goto exit;
		}

		list_link_pop(list_link);

		aoc_debug("ACTIVE (%lu) %li %li %li\n",
			state->dist, state->pos.x, state->pos.y, state->pos.z);

		hmap_link = hmap_get(&distmap, (struct hmap_key) {.p = &state->pos});
		assert(hmap_link);
		if (state->dist != hmap_link->value.s) {
			free(state);
			continue;
		}

		hmap_link = hmap_get(portals, (struct hmap_key) {.p = &state->pos});
		if (hmap_link) {
			id = hmap_link->value.s;
			if (get_portal(portals, &npos, id ^ PORTAL_IN_MASK)) {
				if (id) {
					aoc_debug("PORTAL %c %c %i -> %li %li\n",
						(id >> 8) & 0xff, id & 0xff,
						!!(id & PORTAL_IN_MASK), npos.x, npos.y);
				} else {
					aoc_debug("NO PORTAL %c %c\n",
						(hmap_link->value.s >> 8) & 0xff,
						(hmap_link->value.s >> 0) & 0xff);
				}

				if (id & PORTAL_IN_MASK) {
					npos.z = state->pos.z + 1;
					add_djk_state(&active, &distmap, portals,
						&npos, state->dist + 1);
				} else if (id) {
					npos.z = state->pos.z - 1;
					if (npos.z >= 0) {
						add_djk_state(&active, &distmap, portals,
							&npos, state->dist + 1);
					}
				}
			}
		}

		for (i = 0; i < 4; i++) {
			vec3i_add_2d(&npos, &state->pos, &adj[i]);
			hmap_link = hmap_get(map, (struct hmap_key) {.p = &npos});
			if (!hmap_link || hmap_link->value.c != '.') continue;
			add_djk_state(&active, &distmap, portals,
				&npos, state->dist + 1);
		}

		free(state);
	}

exit:
	for (HMAP_ITER(&distmap, hmap_iter))
		free(hmap_iter.link->key._p);
	hmap_deinit(&distmap);
	list_free_items(&active, free, LIST_OFFSET(struct djk_state, link));

	return mdist;
}

void
part2(void)
{
	struct hmap map;
	struct hmap portals;
	struct hmap_iter iter;
	struct vec3i start, end;
	struct vec2i min, max;
	size_t answer;

	hmap_init(&map, 256, vec2i_hmap_hash, vec2i_hmap_keycmp,
		&stdlib_strict_heap_allocator);
	hmap_init(&portals, 256, vec2i_hmap_hash, vec2i_hmap_keycmp,
		&stdlib_strict_heap_allocator);

	load_map(&map, aoc.input, aoc.input_size, &min, &max);
	load_portals(&portals, &map, &min, &max);

	assert(get_portal(&portals, &start, PORTAL_ID('A', 'A', 0)));
	assert(get_portal(&portals, &end, PORTAL_ID('Z', 'Z', 0)));
	answer = min_dist(&map, &portals, &start, &end);

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

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