// Copyright (c) 2023, Marvin Borner <dev@marvinborner.de>
// SPDX-License-Identifier: MIT

#include <stdio.h>

#include <term.h>
#include <map.h>
#include <reduce.h>
#include <log.h>
#include <assert.h>

// parent rehashing done by caller
static struct term *shift(struct term *term, size_t level, size_t amount)
{
	debug("shifting %lx\n", term->hash & HASH_MASK);
	if (term->type == VAR) {
		if (level > term->u.var.index)
			return term;
		/* term->refs += 1; // TODO: Kinda hacky */
		return term_rehash_var(term, term->u.var.index + amount);
	} else if (term->type == ABS) {
		struct term *previous = term->u.abs.term;
		struct term *new = shift(term->u.abs.term, level + 1, amount);
		if (previous->hash == new->hash)
			return term; // nothing changed
		/* term->refs += 1; */
		return term_rehash_abs(term, new);
	} else if (term->type == APP) {
		hash_t previous_lhs = term->u.app.lhs->hash;
		hash_t previous_rhs = term->u.app.rhs->hash;
		struct term *lhs = shift(term->u.app.lhs, level, amount);
		struct term *rhs = shift(term->u.app.rhs, level, amount);
		if (previous_lhs == lhs->hash && previous_rhs == rhs->hash)
			return term; // nothing changed
		/* term->refs += 1; */
		return term_rehash_app(term, lhs, rhs);
	}
	fatal("invalid type %d\n", term->type);
}

// TODO: remove or to enum
struct substitution_state {
	int a;
};

// only substitutes (and therefore increments indices and rehashes)
// TODO: formalize (write all rules)
// TODO: fix inc.bruijn
// TODO: fix unbound.bruijn
// TODO: try [([[1 1]] [[1]])]
// TODO: fix reparenting
static struct term *substitute(struct term *term, struct term *substitution,
			       size_t level, struct substitution_state *state)
{
	debug("substituting %lx with %lx\n", term->hash & HASH_MASK,
	      substitution->hash & HASH_MASK);
	if (term->type == VAR) {
		if (level == term->u.var.index) { // replace by shifted term
			state->a = 1;
			struct term *shifted = shift(substitution, 0, level);
			if (shifted->hash != substitution->hash) {
				state->a = 4;
			}
			return shifted;
		} else if (level > term->u.var.index) { // crush unbound
			state->a = 2;
			term_deref(substitution, 1);
			return term;
		} else { // shift idx by -1
			state->a = 3;
			struct term *shifted = shift(term, 0, -1);
			term_deref(substitution, 1);
			return shifted;
		}
	} else if (term->type == ABS) {
		/* struct term *previous = term->u.abs.term; */
		hash_t previous_hash = term->u.abs.term->hash;
		/* term->refs += 1; */
		/* previous->refs += 1; */
		struct term *new = substitute(term->u.abs.term, substitution,
					      level + 1, state);
		if (previous_hash == new->hash)
			return term; // nothing changed
		struct term *rehashed = term_rehash_abs(term, new);
		term_rehash_parents(rehashed);
		if (state->a == 3) {
			/* term_deref_head(previous, 1); */
			/* term_deref_head(previous, 1); */
			/* term->refs -= 1; */
			state->a = 0;
		}
		if (state->a == 1) {
			/* term_deref_head(previous, 1); */
			/* term_deref_head(previous, 1); */
			state->a = 0;
		}
		if (state->a == 4) {
			/* term_deref_head(term, 0); */
			/* term_deref_head(previous, 0); */
			state->a = 0;
		}
		return rehashed;
	} else if (term->type == APP) { // no beta reduction; just substitution
		hash_t previous_lhs = term->u.app.lhs->hash;
		hash_t previous_rhs = term->u.app.rhs->hash;
		struct term *lhs =
			substitute(term->u.app.lhs, substitution, level, state);
		struct term *rhs =
			substitute(term->u.app.rhs, substitution, level, state);
		if (previous_lhs == lhs->hash && previous_rhs == rhs->hash)
			return term; // nothing changed
		struct term *rehashed = term_rehash_app(term, lhs, rhs);
		term_rehash_parents(rehashed);
		return rehashed;
	}
	fatal("invalid type %d\n", term->type);
}

// reduction of application
// ([X] Y) -> X/Y
struct term *reduce(struct term *term)
{
	if (!term_is_beta_redex(term))
		fatal("can't reduce non-beta-redex %d\n", term->type);

	debug("reducing %lx\n", term->hash & HASH_MASK);

	struct substitution_state state = { 0 };
	return substitute(term->u.app.lhs, term->u.app.rhs, -1, &state);
}