//#define _GNU_SOURCE

#include <limits.h>
#include <regex.h>
#include <search.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>


#define arrlen(x) (sizeof(x)/sizeof((x)[0]))


enum operation {
	OP_VAL,
	OP_ADD,
	OP_SUB,
	OP_MUL,
	OP_DIV,
	OP_CMP,
	OP_UNK,
};


enum operation
operation_get(char op)
{
	switch (op) {
	case '+': return OP_ADD;
	case '-': return OP_SUB;
	case '*': return OP_MUL;
	case '/': return OP_DIV;
	case '=': return OP_CMP;
	case '?': return OP_UNK;
	}
	return OP_VAL;
}


#define MONKEY_UNKNOWN LONG_MIN

struct monkey {
	char name[5];
	enum operation op;
	long value;
	char left[5];
	char right[5];
};


struct monkey *
monkey_create_val(char const * name, long value)
{
	struct monkey * m;

	if (!name || strlen(name) != 4)
		return NULL;

	if ((m = malloc(sizeof(struct monkey))) == NULL)
		return NULL;

	snprintf(m->name, sizeof(m->name), "%s", name);
	m->op = OP_VAL;
	m->value = value;
	m->left[0] = '\0';
	m->right[0] = '\0';
	return m;
}


struct monkey *
monkey_create_op(char const * name, enum operation op, char const * left, char const * right)
{
	struct monkey * m;

	if (!name || strlen(name) != 4
			|| op == OP_VAL
			|| !left || strlen(left) != 4
			|| !right || strlen(right) != 4)
		return NULL;

	if ((m = malloc(sizeof(struct monkey))) == NULL)
		return NULL;

	snprintf(m->name, sizeof(m->name), "%s", name);
	m->op = op;
	m->value = 0;
	snprintf(m->left, sizeof(m->left), "%s", left);
	snprintf(m->right, sizeof(m->right), "%s", right);
	return m;
}


int
monkey_cmp(struct monkey const * a, struct monkey const * b)
{
	return strcmp(a->name, b->name);
}


int
monkey_cmp_t(void const * a, void const * b)
{
	return monkey_cmp(a, b);
}


struct monkey *
monkey_find(void * monkey_tree, char const * name)
{
	struct monkey tmp, ** result;

	if (!name || strlen(name) != 4) {
		fprintf(stderr, "Invalid monkey name %s\n", name ? name : "(null)");
		return NULL;
	}

	snprintf(tmp.name, sizeof(tmp.name), "%s", name);
	result = tfind(&tmp, &monkey_tree, monkey_cmp_t);
	return result ? *result : NULL;
}


long
monkey_value(void * monkey_tree, struct monkey * m)
{
	long left, right, v;

	if (!m) {
		fprintf(stderr, "%s: Bad monkey\n", __FUNCTION__);
		return 0;
	}

	if (m->op == OP_VAL)
		return m->value;

	if (m->op == OP_UNK)
		return MONKEY_UNKNOWN;

	left = monkey_value(monkey_tree, monkey_find(monkey_tree, m->left));
	right = monkey_value(monkey_tree, monkey_find(monkey_tree, m->right));
	if (left == MONKEY_UNKNOWN || right == MONKEY_UNKNOWN)
		return MONKEY_UNKNOWN;

	switch (m->op) {
	case OP_CMP:
		// If they match, return that value, otherwise MONKEY_UNKNOWN
		v = left == right ? left : MONKEY_UNKNOWN;
		break;

	case OP_ADD:
		v = left + right;
		break;

	case OP_SUB:
		v = left - right;
		break;

	case OP_MUL:
		v = left * right;
		break;

	case OP_DIV:
		v = left / right;
		break;

	default:
		fprintf(stderr, "Unexpected operation for %s: %d\n", m->name, m->op);
		return 0;
	}

	return v;
}


// Find a monkey with OP_UNK in the tree, fixup its value to match the value passed in, and return that monkey.
struct monkey *
monkey_fixunknown(void * monkey_tree, struct monkey * m, long value)
{
	struct monkey * mleft, * mright, * munk;
	long vleft, vright, v;

	if (!m) {
		fprintf(stderr, "%s: Bad monkey\n", __FUNCTION__);
		return NULL;
	}

	if (m->op == OP_UNK) {
		// Is an unknown value. Set it to the value we passed in and return it!
		m->op = OP_VAL;
		m->value = value;
		return m;
	}

	if (m->op == OP_VAL)
		// Can't fix an existing value.
		return NULL;


	// Binary operation. Get the left and right monkeys/values
	mleft = monkey_find(monkey_tree, m->left);
	mright = monkey_find(monkey_tree, m->right);

	vleft = monkey_value(monkey_tree, mleft);
	vright = monkey_value(monkey_tree, mright);

	if (vleft == MONKEY_UNKNOWN && vright == MONKEY_UNKNOWN)
		// Two unknowns. Can't fix.
		return NULL;

	if (vleft != MONKEY_UNKNOWN && vright != MONKEY_UNKNOWN)
		// No unknowns. Nothing to fix.
		return NULL;

	// Work out which is the unknown, and which is the known
	if (vleft == MONKEY_UNKNOWN) {
		munk = mleft;
		v = vright;
	}
	else {
		munk = mright;
		v = vleft;
	}

	if (m->op == OP_CMP) {
		// It's a comparison. If the passed in value is different from
		// the known value, we can't fix.
		if (value != MONKEY_UNKNOWN && value != v)
			return NULL;
		// OK, fix the unknown side.
		return monkey_fixunknown(monkey_tree, munk, v);
	}

	if (value == MONKEY_UNKNOWN)
		// Trying to fix a binary operation, where the result we need
		// to produce is unknown?
		return NULL;

	// OK, we have a binary op where we know the result (#value) and one of
	// the operands, but not the other. Recursively fix the operand we don't
	// know.
	switch (m->op) {
	case OP_ADD:
		// unknown + known = value => unknown = value - known
		return monkey_fixunknown(monkey_tree, munk, value - v);

	case OP_SUB:
		if (vleft == MONKEY_UNKNOWN)
			// unknown - right = value => unknown = value + right
			return monkey_fixunknown(monkey_tree, munk, value + vright);
		else
			// left - unknown = value => unknown = left - value
			return monkey_fixunknown(monkey_tree, munk, vleft - value);

	case OP_MUL:
		// unknown * known = value => unknown = value / known
		return monkey_fixunknown(monkey_tree, munk, value / v);

	case OP_DIV:
		if (vleft == MONKEY_UNKNOWN)
			// unknown / right = value => unknown = value * right
			return monkey_fixunknown(monkey_tree, munk, value * vright);
		else
			// left / unknown = value => unknown = left - value
			return monkey_fixunknown(monkey_tree, munk, vleft / value);

	default:
		fprintf(stderr, "Unexpected operation for %s: %d\n", m->name, m->op);
		break;
	}

	return NULL;
}


int
main(int argc, char ** argv)
{
	int part = 1, i;
	regex_t reval, reop;
	char buf[BUFSIZ];
	void * monkey_tree = NULL;
	struct monkey * m;


	// Parse command options
	while ((i = getopt(argc, argv, "p:")) != -1) {
		switch (i) {
		case 'p':
			part = atoi(optarg);
			break;

		case -1:
			return -1;
		}
	}

	// Compile regexes
	if (regcomp(&reval, "([[:alpha:]]{4}): (-?[[:digit:]]+)", REG_EXTENDED) != 0
			|| regcomp(&reop, "([[:alpha:]]{4}): ([[:alpha:]]{4}) ([-+*/]) ([[:alpha:]]{4})", REG_EXTENDED) != 0)
	{
		fprintf(stderr, "Bad regex\n");
		return -1;
	}

	// Read file.
	while (fgets(buf, sizeof(buf), stdin)) {
		regmatch_t matches[5];
		struct monkey * m;

		if (regexec(&reval, buf, arrlen(matches), matches, 0) == 0) {
			buf[matches[1].rm_eo] = '\0';
			m = monkey_create_val(buf + matches[1].rm_so, atoi(buf + matches[2].rm_so));
		}
		else if (regexec(&reop, buf, arrlen(matches), matches, 0) == 0) {
			buf[matches[1].rm_eo] = '\0';
			buf[matches[2].rm_eo] = '\0';
			buf[matches[4].rm_eo] = '\0';
			m = monkey_create_op(buf + matches[1].rm_so,
					operation_get(buf[matches[3].rm_so]),
					buf + matches[2].rm_so,
					buf + matches[4].rm_so);
		}
		else {
			fprintf(stderr, "Bad monkey: %s\n", buf);
#ifdef _GNU_SOURCE
			tdestroy(monkey_tree, free);
#endif
			regfree(&reop);
			regfree(&reval);
			return -1;
		}

		if (part == 2 && strcmp(m->name, "root") == 0) {
			m->op = OP_CMP;
			m->value = MONKEY_UNKNOWN;
		}
		else if (part == 2 && strcmp(m->name, "humn") == 0) {
			m->op = OP_UNK;
			m->value = MONKEY_UNKNOWN;
		}

		tsearch(m, &monkey_tree, monkey_cmp_t);
	}

	// Output result.
	m = monkey_find(monkey_tree, "root");
	switch (part) {
	case 1:
		printf("root value: %ld\n", monkey_value(monkey_tree, m));
		break;

	case 2:
		if ((m = monkey_fixunknown(monkey_tree, m, MONKEY_UNKNOWN)) == NULL)
			printf("Unknown value not found\n");
		else
			printf("Fixed unknown value for %s: %ld\n", m->name, m->value);
		break;

	default:
		fprintf(stderr, "Unexpected part %d\n", part);
#ifdef _GNU_SOURCE
		tdestroy(monkey_tree, free);
#endif
		regfree(&reop);
		regfree(&reval);
		return -1;
	}

	// Tidy and done.
#ifdef _GNU_SOURCE
	tdestroy(monkey_tree, free);
#endif
	regfree(&reop);
	regfree(&reval);

	return 0;
}