/* SPDX-License-Identifier: GPL-2.0-only OR GPL-3.0-only */
/* Copyright (c) 2020-2024 Brett Sheffield <bacs@librecast.net> */

#include <assert.h>
#include <librecast_pvt.h>
#include <librecast/crypto.h>
#include <librecast/mtree.h>
#include <math.h>
#include <pthread.h>
#include <semaphore.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>

struct mtree_job_s {
	mtree_t *tree;
	sem_t *node;  /* semaphores for parent nodes above subtree root level */
	size_t root;  /* root node of subtree */
	int subtrees; /* tree is divided into how many subtrees */
};

/* next_pow2 - Public Domain, credit to Sean Anderson from
 * https://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2 */
static uint32_t next_pow2(uint32_t v)
{
	v--;
	v |= v >> 1;
	v |= v >> 2;
	v |= v >> 4;
	v |= v >> 8;
	v |= v >> 16;
	return ++v;
}

void mtree_hexdump(mtree_t *tree, FILE *fd)
{
	char hex[HEXLEN];
	for (size_t i = 0; i < tree->nodes; i++) {
		sodium_bin2hex(hex, HEXLEN, tree->tree + i * HASHSIZE, HASHSIZE);
		fprintf(fd, "%08zu: %.*s\n", i, HEXLEN, hex);
	}
}

void mtree_printmap(unsigned char *map, size_t len, FILE *fd)
{
	for (size_t i = 0; i < len; i++) {
		fprintf(fd, "%d", !!isset(map, i));
	}
	fputc('\n', fd);
}

size_t mtree_child(mtree_t *tree, size_t node)
{
	node = (node + 1) * 2 - 1;
	return (node >= tree->nodes) ? 0 : node;
}

uint8_t *mtree_nnode(mtree_t *t, size_t node)
{
	if (node >= t->nodes) return (errno = ERANGE), NULL;
	return t->tree + node * MTREE_HASHSIZE;
}

static size_t mtree_base(mtree_t *tree)
{
	return next_pow2(tree->chunks);
}

/* return number of levels for tree with base number of data nodes */
static size_t mtree_levels(size_t base)
{
	return (size_t)log2(next_pow2((uint64_t)base)) + 1;
}

static size_t mtree_node_level(size_t node)
{
	return (size_t)log2(node + 1);
}

/* node numbered from 0=root, levels numbered from 0=base */
static size_t mtree_node_level_base(size_t base, size_t node)
{
	size_t rootlvl = mtree_node_level(node);
	size_t treelvl = mtree_levels(base);
	assert(treelvl >= rootlvl);
	return treelvl - rootlvl - 1;
}

/* return number of nodes in tree with base number of data nodes */
static size_t mtree_size(size_t base)
{
	return (size_t)(next_pow2((uint32_t)base) << 1) - 1;
}

static size_t mtree_child_base(size_t base, size_t node)
{
	node = (node + 1) * 2 - 1;
	return (node >= mtree_size(base)) ? 0 : node;
}

size_t mtree_subtree_data_max(size_t base, size_t root)
{
	size_t n = root;
	while (mtree_node_level_base(base, n)) {
		n = mtree_child_base(base, n) + 1;
	}
	return n;
}

size_t mtree_subtree_data_min(size_t base, size_t root)
{
	size_t n = root;
	while (mtree_node_level_base(base, n)) {
		n = mtree_child_base(base, n);
	}
	return n;
}

static size_t mtree_blocks_subtree(mtree_t *tree, size_t root)
{
	size_t base = mtree_base(tree);
	size_t max = mtree_subtree_data_max(base, root);
	size_t min = mtree_subtree_data_min(base, root);
	size_t tmin = mtree_subtree_data_min(base, 0);
	max = MIN(max, tmin + tree->chunks - 1);
	return max - min + 1;
}

static size_t mtree_node_offset(size_t node)
{
	size_t npow = 1;
	size_t lvl = mtree_node_level(node);
	while (lvl--) npow *= 2;
	return node - npow + 1;
}

/* NB: we have no bounds checking or errors,
 * root MUST NOT be greater than node */
static size_t mtree_node_offset_subtree(size_t node, size_t root)
{
	size_t off = mtree_node_offset(node);
	size_t a = mtree_node_level(node);
	size_t b = mtree_node_level(root);
	size_t d = (1U << a) / (1U << b);
	while (off > (d - 1)) off -= d;
	return off;
}

static size_t mtree_chunkn_len(mtree_t *tree, size_t n)
{
	size_t mod;
	size_t min = mtree_subtree_data_min(mtree_base(tree), 0);
	size_t max = mtree_subtree_data_max(mtree_base(tree), 0);
	max = MIN(max, min + tree->chunks - 1);
	if (n < min || n > max || n > min + tree->chunks) return 0;
	mod = tree->len % MTREE_CHUNKSIZE;
	return ((mod) && (n == max)) ? mod : MTREE_CHUNKSIZE;
}

/* set map bits for leaf node */
static unsigned char *leafnode(unsigned char *map, mtree_t *t1,
		size_t root, size_t node, unsigned bits, int set_node)
{
	size_t blen = mtree_chunkn_len(t1, node);
	unsigned bitcount = howmany(blen, MTREE_CHUNKSIZE / bits);
	if (set_node) node = mtree_node_offset_subtree(node, root);
	for (unsigned i = 0; i < bitcount; i++) {
		setbit(map, node * bits + i);
	}
	return map;
}

static unsigned char *mtree_search_subtree(unsigned char *map, mtree_t *t1, mtree_t *t2,
		size_t root, size_t node, size_t child, unsigned bits)
{
	size_t *q; /* FIFO queue */
	int i = 0, o = 0; /* in / out */
	q = calloc(t1->nodes, sizeof(size_t));
	if (!q) return NULL;
	q[i++] = child; child++;
	q[i++] = child;
	while ((node = q[o++])) {
		if (memcmp(mtree_nnode(t1, node), mtree_nnode(t2, node), HASHSIZE)) {
			child = mtree_child(t1, node);
			if (child) {
				q[i++] = child; child++;
				q[i++] = child;
			}
			else leafnode(map, t1, root, node, bits, 1);
		}
	}
	free(q);
	return map;
}

/* perform bredth-first search of subtree, return bitmap
 * bits = number of bits to use per chunk. This is useful if chunksize > MTU, a chunk will be
 * fragmented into multiple packets
 * map must be free()d by caller */
unsigned char *mtree_diff_subtree(mtree_t *t1, mtree_t *t2, size_t root, unsigned bits)
{
	unsigned char *map = NULL;
	size_t base, child, node;
	if (!memcmp(mtree_nnode(t1, root), mtree_nnode(t2, root), HASHSIZE))
		return NULL; /* subtree root matches, stop now */

	base = mtree_blocks_subtree(t1, root);
	if (!base) return NULL; /* zero blocks */
	node = (base + (CHAR_BIT - 1)) / CHAR_BIT;
	map = calloc(bits, base);
	if (!map) return NULL;
	child = mtree_child(t1, root);
	if (!child) {
		node = mtree_node_offset_subtree(node, root);
		return leafnode(map, t1, root, node, bits, 0);
	}
	return mtree_search_subtree(map, t1, t2, root, node, child, bits);
}

unsigned char *mtree_diff_map(mtree_t *t1, mtree_t *t2)
{
	return mtree_diff_subtree(t1, t2, 0, 1);
}

int mtree_verify(mtree_t *t)
{
	unsigned char hash[HASHSIZE];
	if (!t) return -1;
	if (!t->len) return 0;
	if (t->tree == NULL) return -1;
	for (size_t node = 1; node < t->nodes - 1; node += 2) {
		hash_generic(hash, HASHSIZE, mtree_nnode(t, node), MTREE_HASHSIZE * 2);
		if (memcmp(hash, mtree_nnode(t, mtree_parent(node)), HASHSIZE))
			return (errno = EBADMSG), -1;
	}
	return 0;
}

static void mtree_hashparents(struct mtree_job_s *job, size_t node, size_t root)
{
	size_t left = node - 1;
	size_t parent = mtree_parent(node);
	uint8_t *hash =  job->tree->tree + parent * MTREE_HASHSIZE;
	unsigned char *in = job->tree->tree + left * MTREE_HASHSIZE;
	if (left < root) sem_wait(&job->node[left]);
	hash_generic(hash, HASHSIZE, in, MTREE_HASHSIZE * 2);
	if (parent <= root) sem_post(&job->node[parent]);
	if (parent && !(parent & 1)) mtree_hashparents(job, parent, root);
}

static void * mtree_buildsubtree(void *arg)
{
	struct mtree_job_s *job = (struct mtree_job_s *)arg;
	/* atomically find root for our subtree */
	size_t root = fetch_add(&job->root, 1);
	size_t noff = mtree_node_offset(root);
	size_t min = noff * job->tree->base / job->subtrees;
	size_t max = min + job->tree->base / job->subtrees - 1;
	for (size_t chunk = min; chunk <= max; chunk++) {
		/* node => chunk number (base node), starting from 0 */
		size_t node = chunk + job->tree->nodes - job->tree->base;
		if (chunk < job->tree->chunks) {
			size_t off = chunk * MTREE_CHUNKSIZE;
			size_t len = MIN(MTREE_CHUNKSIZE, job->tree->len - off);
			uint8_t *hash =  job->tree->tree + node * MTREE_HASHSIZE;
			unsigned char *in = (unsigned char *)job->tree->data + off;
			hash_generic(hash, HASHSIZE, in, len);
		}
		if (node <= root) sem_post(&job->node[node]);
		if (chunk & 1) mtree_hashparents(job, node, root);
	}
	return arg;
}

static int mtree_hashsubtrees(mtree_t *t, q_t *q, int nthreads)
{
	struct mtree_job_s job = { .tree = t };
	const int subtrees = MIN(nthreads, (int)t->base);
	const int semcount = subtrees * 2 - 1;
	sem_t node[semcount];
	int rc = -1;
	job.root = subtrees - 1; /* root of first subtree */
	job.node = node;
	job.subtrees = subtrees;
	for (int i = 0; i < semcount; i++) sem_init(&node[i], 0, 0);
	for (int i = 0; i < subtrees; i++) {
		q_push(q, mtree_buildsubtree, &job);
	}
	rc = sem_wait(&node[0]);
	for (int i = 0; i < semcount; i++) sem_destroy(&node[i]);
	return rc;
}

static int hash_subtrees_queue(mtree_t *t, int nthreads)
{
	pthread_t tid[nthreads];
	q_t q = {0};
	int rc;
	q_init(&q);
	q_pool_create(tid, nthreads, q_job_seek, &q);
	rc = mtree_hashsubtrees(t, &q, nthreads);
	q_pool_destroy(tid, nthreads);
	q_free(&q);
	return rc;
}

int mtree_build(mtree_t *t, void * const data, q_t *q)
{
	int nthreads = MIN(MTREE_THREADMAX, t->base);
	if (!t->len) return 0;
	t->data = data;
	if (!q) return hash_subtrees_queue(t, nthreads);
	else return mtree_hashsubtrees(t, q, nthreads);
}

int mtree_init(mtree_t *t, size_t sz)
{
	memset(t, 0, sizeof *t);
	t->len = sz;
	t->chunks = MAX(howmany(sz, MTREE_CHUNKSIZE), 1);
	t->base = next_pow2(t->chunks);
	t->nodes = (t->base << 1) - 1;
	t->tree = calloc(t->nodes, MTREE_HASHSIZE);
	return (t->tree) ? 0 : -1;
}

void mtree_free(mtree_t *t)
{
	int err = errno;
	free(t->tree);
	t->data = NULL;
	errno = err;
}
