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

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

struct ingredient {
	char name[8];
	size_t amount;
};

struct reaction {
	struct dvec ingredients;
	struct ingredient result;

};

void
parse_reactions(struct hmap *reactions)
{
	const char *pos, *sep, *lpos, *lend;
	struct reaction *reaction;
	struct ingredient ingredient;
	char line[256];
	void *slot;
	int rc;

	lpos = aoc.input;
	lend = lpos + aoc.input_size;
	while (readtok(line, sizeof(line), '\n', &lpos, lend)) {
		if (strspn(line, " \r\n\t\v") == strlen(line))
			continue;

		reaction = malloc(sizeof(struct reaction));
		assert(reaction != NULL);
		dvec_init(&reaction->ingredients, sizeof(struct ingredient),
			10, &stdlib_strict_heap_allocator);

		pos = line;
		while (1) {
			sep = strchr(pos, ',');
			if (!sep) sep = strchr(pos, '>');

			rc = sscanf(pos, "%lu %[^,= ]",
				&ingredient.amount, ingredient.name);
			assert(rc == 2);

			if (sep) {
				slot = dvec_add_slot(&reaction->ingredients);
				memcpy(slot, &ingredient,
					sizeof(struct ingredient));
			} else {
				memcpy(&reaction->result, &ingredient,
					sizeof(struct ingredient));
			}

			if (!sep) break;
			pos = sep + 2;
		}

		hmap_add(reactions,
			(struct hmap_key) {.p = reaction->result.name},
			(struct hmap_val) {.p = reaction});
	}
}

size_t
add_excess(struct hmap *hmap, struct ingredient *result, struct reaction *reaction)
{
	struct hmap_link **link;
	size_t n, amount, excess;

	amount = result->amount;
	link = hmap_link_pos(hmap, (struct hmap_key) {.p = reaction->result.name});
	if (link && *link) {
		aoc_debug("get excess %s = %u\n", reaction->result.name,
			(*link)->value.u);
		if ((*link)->value.u >= amount) {
			(*link)->value.u -= (unsigned int) amount;
			return 0;
		} else {
			amount -= (*link)->value.u;
			(*link)->value.u = 0;
		}
	}

	n = CEILDIV(amount, reaction->result.amount);
	excess = n * reaction->result.amount - amount;
	aoc_debug("set excess %s = %lu\n", reaction->result.name, excess);
	if (link && *link) {
		(*link)->value.u = (unsigned int) excess;
	} else {
		hmap_add(hmap, (struct hmap_key) {.p = reaction->result.name},
			(struct hmap_val) {.u = (unsigned int) excess});
	}

	return n;
}

size_t
calculate_ore(struct hmap *reactions, struct hmap *excess,
	const char *name, size_t amount)
{
	struct hmap_link *link;
	struct reaction *reaction;
	struct ingredient ingredient;
	struct ingredient result;
	struct dvec reactants;
	struct ingredient *slot, *slot2;
	size_t ore, n;

	dvec_init(&reactants, sizeof(struct ingredient),
		10, &stdlib_strict_heap_allocator);

	slot = dvec_add_slot(&reactants);
	ingredient.amount = amount;
	snprintf(ingredient.name, 8, "%s", name);
	memcpy(slot, &ingredient, sizeof(struct ingredient));

	ore = 0;
	while (!dvec_empty(&reactants)) {
		slot = dvec_back(&reactants);
		memcpy(&result, slot, sizeof(struct ingredient));
		dvec_rm_slot(&reactants, slot);

		link = hmap_get(reactions, (struct hmap_key) {.p = result.name});
		assert(link != NULL);
		reaction = link->value._p;

		n = add_excess(excess, &result, reaction);
		if (!n) continue;

		for (DVEC_ITER(&reaction->ingredients, slot)) {
			aoc_debug("%lu %s needs %lu %s\n",
				result.amount, result.name,
				slot->amount * n, slot->name);
			if (!strcmp(slot->name, "ORE")) {
				ore += slot->amount * n;
			} else {
				snprintf(ingredient.name, 8, "%s", slot->name);
				ingredient.amount = slot->amount * n;
				slot2 = dvec_add_slot(&reactants);
				memcpy(slot2, &ingredient, sizeof(struct ingredient));
			}
		}
	}

	dvec_deinit(&reactants);

	return ore;
}

void
part1(void)
{
	struct hmap_iter iter;
	struct hmap reactions;
	struct hmap excess;
	struct reaction *reaction;
	size_t ore;

	hmap_init(&reactions, 256, hmap_str_hash, hmap_str_keycmp,
		&stdlib_strict_heap_allocator);
	hmap_init(&excess, 256, hmap_str_hash, hmap_str_keycmp,
		&stdlib_strict_heap_allocator);

	parse_reactions(&reactions);

	ore = calculate_ore(&reactions, &excess, "FUEL", 1);

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

	for (HMAP_ITER(&reactions, iter)) {
		reaction = iter.link->value._p;
		dvec_deinit(&reaction->ingredients);
		free(iter.link->value._p);
	}
	hmap_deinit(&excess);
	hmap_deinit(&reactions);
}

void
part2(void)
{
	const size_t max_ore = 1000000000000UL;
	struct hmap_iter iter;
	struct hmap reactions;
	struct hmap excess;
	struct reaction *reaction;
	size_t fuel_min, fuel_max;
	size_t ore, fuel;

	hmap_init(&reactions, 256, hmap_str_hash, hmap_str_keycmp,
		&stdlib_strict_heap_allocator);
	hmap_init(&excess, 256, hmap_str_hash, hmap_str_keycmp,
		&stdlib_strict_heap_allocator);

	parse_reactions(&reactions);

	fuel_min = 1;
	fuel_max = max_ore;
	while (1) {
		fuel = (fuel_min + fuel_max) / 2;
		if (fuel == fuel_min) break;
		ore = calculate_ore(&reactions, &excess, "FUEL", fuel);
		if (ore > max_ore) {
			fuel_max = fuel;
		} else if (ore < max_ore) {
			fuel_min = fuel;
		} else {
			break;
		}
		hmap_clear(&excess);
	}

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

	for (HMAP_ITER(&reactions, iter)) {
		reaction = iter.link->value._p;
		dvec_deinit(&reaction->ingredients);
		free(iter.link->value._p);
	}
	hmap_deinit(&excess);
	hmap_deinit(&reactions);
}