#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) ((size_t) ((uint8_t) (a) * 100 + (uint8_t) (b)))

struct djk_state {
	size_t dist;
	struct vec2i pos;

	struct list_link link;
};

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)
{
	char line[256];
	const char *lpos, *lend;
	struct vec2i pos, *key;

	lpos = aoc.input;
	lend = lpos + aoc.input_size;
	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;
			key = malloc(sizeof(struct vec2i));
			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 hmap_link *link1;
	struct hmap_link *link2;
	struct hmap_iter iter;
	struct vec2i pos;
	struct vec2i npos;
	struct vec2i *key;
	size_t i, id;

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

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

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

		aoc_debug("FOUND %c %c\n", iter.link->value.c, link1->value.c);

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

static bool
get_portal(struct hmap *portals, struct vec2i *pos, struct vec2i *skip, size_t id)
{
	struct hmap_iter iter;

	for (HMAP_ITER(portals, iter)) {
		if (skip && vec2i_eql(skip, iter.link->key.p))
			continue;
		if (iter.link->value.s == id) {
			vec2i_set(pos, iter.link->key.p);
			return true;
		}
	}

	return false;
}

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

	hmap_linkp = hmap_link_pos(distmap, (struct hmap_key) {.p = pos});
	if (!*hmap_linkp) {
		key = memdup(pos, sizeof(struct vec2i));
		*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) {
		nstate = malloc(sizeof(struct djk_state));
		nstate->dist = dist;
		vec2i_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 vec2i *start, struct vec2i *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 vec2i *key, npos;
	size_t mdist, i;

	list_init(&active);
	hmap_init(&distmap, 256, vec2i_hmap_hash, vec2i_hmap_keycmp,
		&stdlib_strict_heap_allocator);

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

	key = memdup(start, sizeof(struct vec2i));
	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 (vec2i_eql(&state->pos, end)) {
			mdist = state->dist;
			goto exit;
		}

		list_link_pop(list_link);

		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) {
			if (get_portal(portals, &npos, &state->pos, hmap_link->value.s))
				add_djk_state(&active, &distmap, portals,
					&npos, state->dist + 1);
		}

		for (i = 0; i < 4; i++) {
			vec2i_add(&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
part1(void)
{
	struct hmap map;
	struct hmap portals;
	struct hmap_iter iter;
	struct vec2i start, end;
	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);
	load_portals(&portals, &map);

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

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

	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);
}