// 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 <term.h>
#include <map.h>

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, 0);
	return term;
}

void term_print(struct term *term)
{
	switch (term->type) {
	case ABS:
		fprintf(stderr, "[");
		term_print(term->u.abs.term);
		fprintf(stderr, "]");
		break;
	case APP:
		fprintf(stderr, "(");
		term_print(term->u.app.lhs);
		fprintf(stderr, " ");
		term_print(term->u.app.rhs);
		fprintf(stderr, ")");
		break;
	case VAR:
		fprintf(stderr, "%ld", term->u.var.index);
		break;
	default:
		fatal("invalid type %d\n", term->type);
	}
}

struct term *term_rehash_abs(struct term *head, struct term *term)
{
	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
		return match;
	} else { // create new
		struct term *new = term_new(ABS, res, head->depth);
		new->u.abs.term = term;
		map_set(map_all_terms(), new);
		map_set(term->parents, new);
		return new;
	}
}

struct term *term_rehash_app(struct term *head, struct term *lhs,
			     struct term *rhs)
{
	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
		return match;
	} else { // create new
		struct term *new = term_new(APP, res, head->depth);
		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);
		return new;
	}
}

struct term *term_rehash_var(struct term *head, size_t 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
		return match;
	} else { // create new
		struct term *new = term_new(APP, res, head->depth);
		new->u.var.index = index;
		map_set(map_all_terms(), new);
		return new;
	}
}

struct term *term_rehash(struct term *term)
{
	if (term->type == ABS) {
		return term_rehash_abs(term, term->u.abs.term);
	} else if (term->type == APP) {
		return term_rehash_app(term, term->u.app.lhs, term->u.app.rhs);
	} else if (term->type == VAR) {
		return term_rehash_var(term, term->u.var.index);
	}
	return term;
}

// returns the direct parent
void term_rehash_parents(struct term *term)
{
	if (!term->parents)
		return;

	// 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);
		if (previous == new->hash) {
			struct parent_list *next = iterator->next;
			free(iterator);
			iterator = next;
		}

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

		term_rehash_parents(new);

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

void term_refer_head(struct term *term, size_t depth)
{
	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);
}

void term_deref_head(struct term *term)
{
	term->refs--;
	if (term->refs == 0) {
		map_delete(map_all_terms(), term);
		map_destroy(term->parents);
		free(term);
	}
}

void term_deref(struct term *term)
{
	if (term->type == ABS) {
		term_deref(term->u.abs.term);
	} else if (term->type == APP) {
		term_deref(term->u.app.lhs);
		term_deref(term->u.app.rhs);
	}

	term_deref_head(term);
}