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

#include <stdlib.h>
#include <stdio.h>
#include <assert.h>

#include <log.h>
#include <schedule.h>
#include <term.h>
#include <map.h>

static void term_destroy_parent(void *item)
{
	struct term *term = *(struct term **)item;
	term_destroy_head(term, 1);
}

// doesn't care about ref count
// this is needed to destroy possibly multiple parents
void term_destroy_head(struct term *term, char including_parents)
{
	debug("destroying head of %lx\n", term->hash & HASH_MASK);

	map_delete(map_all_terms(), term);

	// recursively destroy own parents
	if (including_parents) {
		map_destroy(term->parents);

		// remove head from child's parents
		if (term->type == ABS) {
			map_delete(term->u.abs.term->parents, term);
		} else if (term->type == APP) {
			map_delete(term->u.app.lhs->parents, term);
			map_delete(term->u.app.rhs->parents, term);
		}
	}

	schedule_remove(term);
	free(term);
}

struct term *term_new(term_type_t type, hash_t hash, size_t depth)
{
	struct term *term = malloc(sizeof(*term));
	if (!term)
		fatal("out of memory!\n");
	term->type = type;
	term->refs = 1;
	term->hash = hash;
	term->depth = depth;
	term->parents =
		hashmap_new(sizeof(struct term *), 0, term_destroy_parent);
	return term;
}

char term_is_beta_redex(struct term *term)
{
	return term->type == APP && term->u.app.lhs->type == ABS;
}

size_t term_print(struct term *term)
{
	switch (term->type) {
	case ABS:
		fprintf(stderr, "[");
		size_t t = term_print(term->u.abs.term);
		fprintf(stderr, "]");
		return 2 + t;
	case APP:
		fprintf(stderr, "(");
		size_t l = term_print(term->u.app.lhs);
		fprintf(stderr, " ");
		size_t r = term_print(term->u.app.rhs);
		fprintf(stderr, ")");
		return 3 + l + r;
	case VAR:
		return fprintf(stderr, "%ld", term->u.var.index);
	default:
		fatal("invalid type %d\n", term->type);
	}
}

struct term *term_rehash_abs(struct term *head, struct term *term,
			     int including_parents)
{
	debug("rehashing abs %lx (%lx)\n", head->hash & HASH_MASK,
	      term->hash & HASH_MASK);
	hash_t res =
		hash((uint8_t *)&head->type, sizeof(head->type), term->hash);

	if (res == head->hash)
		return head;

	struct term *match = map_get(map_all_terms(), res);
	if (match) { // already exists
		// TODO: something different if match->u.abs.term == term)?
		/* term_refer_head(match, head->depth); */
		/* term_deref_head(head, including_parents); */
		return match;
	} else { // create new
		struct term *new = term_new(ABS, res, head->depth);
		// TODO: Clone parents
		new->u.abs.term = term;
		map_set(map_all_terms(), new);
		map_set(term->parents, new);
		/* term_deref_head(head, including_parents); */
		return new;
	}
}

// TODO: main problem is reparenting substitution tickle-up
// TODO: investigate reparenting using gen.blc? ?!!?!
struct term *term_rehash_app(struct term *head, struct term *lhs,
			     struct term *rhs, int including_parents)
{
	debug("rehashing app %lx (%lx, %lx)\n", head->hash & HASH_MASK,
	      lhs->hash & HASH_MASK, rhs->hash & HASH_MASK);
	hash_t res =
		hash((uint8_t *)&head->type, sizeof(head->type), lhs->hash);
	res = hash((uint8_t *)&res, sizeof(res), rhs->hash);

	if (res == head->hash)
		return head;

	struct term *match = map_get(map_all_terms(), res);
	if (match) { // already exists
		/* term_refer_head(match, head->depth); */
		/* term_deref(head, including_parents); */
		return match;
	} else { // create new
		struct term *new = term_new(APP, res, head->depth);
		// TODO: Clone parents
		new->u.app.lhs = lhs;
		new->u.app.rhs = rhs;
		map_set(map_all_terms(), new);
		map_set(lhs->parents, new);
		map_set(rhs->parents, new);
		/* term_deref_head(head, including_parents); */
		return new;
	}
}

struct term *term_rehash_var(struct term *head, size_t index,
			     int including_parents)
{
	debug("rehashing var %lx (%lu)\n", head->hash & HASH_MASK, index);
	hash_t res = hash((uint8_t *)&head->type, sizeof(head->type), index);

	if (res == head->hash)
		return head;

	struct term *match = map_get(map_all_terms(), res);
	if (match) { // already exists
		term_refer_head(match, head->depth);
		/* term_deref(head, including_parents); */
		return match;
	} else { // create new
		struct term *new = term_new(VAR, res, head->depth);
		// TODO: Clone parents
		new->u.var.index = index;
		map_set(map_all_terms(), new);
		/* term_deref_head(head, including_parents); */
		return new;
	}
}

struct term *term_rehash(struct term *term, int including_parents)
{
	if (term->type == ABS)
		return term_rehash_abs(term, term->u.abs.term,
				       including_parents);

	if (term->type == APP)
		return term_rehash_app(term, term->u.app.lhs, term->u.app.rhs,
				       including_parents);

	if (term->type == VAR)
		return term_rehash_var(term, term->u.var.index,
				       including_parents);

	fatal("invalid type %d\n", term->type);
}

void term_rehash_parents(struct term *term)
{
	if (!term->parents)
		return;
	debug("rehashing parents of %lx\n", term->hash & HASH_MASK);

	// we need to convert the parent hashmap to a list
	// so we can replace the rehashed elements while looping
	// TODO: Abstract list lib?
	struct parent_list {
		struct term *term;
		struct parent_list *next;
	};
	struct parent_list *parents = calloc(sizeof(*parents), 1);

	size_t iter = 0;
	void *iter_val;
	while (hashmap_iter(term->parents, &iter, &iter_val)) {
		struct parent_list *new = malloc(sizeof(*parents));
		new->term = *(struct term **)iter_val;
		new->next = parents;
		parents = new;
	}

	struct parent_list *iterator = parents;
	while (iterator && iterator->term) {
		struct term *parent = iterator->term;
		hash_t previous = parent->hash;
		struct term *new = term_rehash(parent, 1);
		if (previous == new->hash) {
			struct parent_list *next = iterator->next;
			free(iterator);
			iterator = next;
			continue;
		}

		map_delete(term->parents, parent);
		map_set(term->parents, new);

		term_rehash_parents(new);

		struct parent_list *next = iterator->next;
		free(iterator);
		iterator = next;
	}
	if (iterator == parents)
		free(parents);
}

void term_refer_head(struct term *term, size_t depth)
{
	debug("referring head of %lx\n", term->hash & HASH_MASK);
	term->refs++;
	if (depth < term->depth) // lower depths are more important
		term->depth = depth;
}

void term_refer(struct term *term, size_t depth)
{
	if (term->type == ABS) {
		term_refer(term->u.abs.term, depth + 1);
	} else if (term->type == APP) {
		term_refer(term->u.app.lhs, depth + 1);
		term_refer(term->u.app.rhs, depth + 1);
	}

	term_refer_head(term, depth);
}

char term_deref_head(struct term *term, char destroy_parents)
{
	debug("dereferring head of %lx\n", term->hash & HASH_MASK);

	assert(term->refs > 0);
	term->refs--;
	if (!term->refs) {
		term_destroy_head(term, destroy_parents);
		return 1;
	}
	return 0;
}

// returns 1 if destroyed
char term_deref(struct term *term, char destroy_parents)
{
	char a = 0, b = 0;
	if (term->type == ABS) {
		a = term_deref(term->u.abs.term, destroy_parents);
	} else if (term->type == APP) {
		a = term_deref(term->u.app.lhs, destroy_parents);
		b = term_deref(term->u.app.rhs, destroy_parents);
	}

	if (a || b)
		return 1;
	return term_deref_head(term, destroy_parents);
}