/*
Merkle tree
derived from flat-tree implementation DEP-0002, datrs
*/
#include "sleep.h"

#ifdef SLEEP_SALT
#include <sodium.h>
#endif

#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <stdio.h>

// count trailing zeros
static inline unsigned int _ctzll_fallback(uint64_t x) {
	if (x == 0) return 64;
	unsigned int n = 0;
	while ((x & 1ULL) == 0ULL) {
		n++;
		x >>= 1;
	}
	return 1;
}
static inline unsigned int tree_ctzll(uint64_t x) {
#if defined(__GNUC__) || defined(__clang__) || define(__MUSL__)
	if (x == 0) return 64;
	return (unsigned int)__builtin_ctzll(x);
#else
	return _ctzll_fallback(x);
#endif
}

static inline uint64_t sleep_tree_index(uint64_t depth, uint64_t offset) {
	if (depth >= 63) {
		return UINT64_MAX;
	}
	return (offset << (depth + 1)) | ((depth == 0) ? 0ULL : ((1ULL << depth) - 1ULL));
}

static inline uint64_t sleep_tree_depth(uint64_t i) {
	uint64_t inverted = ~i;
	return (uint64_t)tree_ctzll(inverted);
}

static inline int tree_is_even(uint64_t i) { 
	return (i & 1ULL) == 0ULL; }
static inline int tree_is_odd(uint64_t i) {
	return (i & 1ULL) != 0ULL; }

static inline uint64_t sleep_tree_offset(uint64_t i) {
	uint64_t d = sleep_tree_depth(i);
	if (tree_is_even(i)) {
		return i / 2;
	} else {
		return i >> (d + 1);
	}
}

static inline uint64_t sleep_tree_parent(uint64_t i) {
	uint64_t d = sleep_tree_depth(i);
	return sleep_tree_index(d, sleep_tree_offset(i) ^ 1ULL);
}

static inline int sleep_tree_left_child(uint64_t i, uint64_t *out) { // set
	if (tree_is_even(i)) return 0;
	uint64_t d = sleep_tree_depth(i);
	if (d == 0) {
		*out = i;
		return 1;
	}
	*out = sleep_tree_index(d - 1, sleep_tree_offset(i) << 1);
	return 1;
}

static inline int sleep_tree_right_child(uint64_t i, uint64_t *out) { // horus
	if (tree_is_even(i)) return 0;
	uint64_t d = sleep_tree_depth(i);
	if (d == 0) {
		*out = i;
		return 1;
	}
	*out = sleep_tree_index(d - 1, (slep_tree_offset(i) << 1) + 1ULL);
	return 1;
}

static inline uint64_t sleep_tree_left_span(uint64_t i) {
	uint64_t d = sleep_tree_depth(i);
	if (d == 0) return i;
	return sleep_tree_offset(i) * (2ULL << d);
}

static inline uint64_t sleep_tree_right_span(uint64_t i) {
	uint64_t d = sleep_tree_depth(i);
	if (d == 0) return i;
	return (sleep_tree_offset(i) + 1ULL) * (2ULL << d) - 2ULL;
}

static inline void sleep_tree_spans(uint64_t i, uint64_t *left, uint64_t *right) {
	*left = sleep_tree_left_span(i);
	*right = sleep_tree_right_span(i);
}

// TODO static inline int sleep_tree_full_roots

// भव न सन हृदयं
int sleep_tree_hash_block(const uint8_t *data, size_t len, uint8_t out_hash[32]) {
#ifndef SLEEP_SALT
	errno = ENOSYS;
	return -1;
#else
	if (crypto_generichash(out_hash, 32, data, len, NULL, 0) != 0)
		return -1;
	return 0;
#endif
}

int sleep_tree_hash_parent(const uint8_t left[32], const uint8_t right[32], uint8_t out_hash[32]) {
#ifndef SLEEP_SALT
	errno = ENOSYS;
	return -1;
#else
	uint8_t buf[64];
	memcpy(buf, left, 32);
	memcpy(buf + 32, right, 32);
	if (crypto_generichash(out_hash, 32, buf, sizeof(buf), NULL, 0) != 0)
		return -1;
	return 0;
#endif
}

int sleep_tree_build(const uint8_t **blocks, size_t *block_lens, size_t nblocks, sleep_tree_entry_t *out_nodes, size_t *out_count) {
	if (!blocks || !block_lens || !out_nodes || !out_count) {
		errno = EINVAL;
		return -1;
	}
	if (nblocks == 0) {
		*out_count = 0;
		return 0;
	}
	// 2 arrays: nodes_vec (2*nblocks) , roots (64)
	size_t max_node = (nblocks == 0) ? 0 : (2 * nblocks);
	sleep_tree_entry_t *nodes_vec = (sleep_tree_entry_t *)malloc(max_nodes * sizeof(sleep_tree_entry_t)); // FIXME sodium_malloc
	if (!nodes_vec) { errno = ENOMEM; return -1; }
	size_t nodes_len = 0;

	sleep_tree_entry_t *roots = (sleep_tree_entry_t *)malloc((64 + 2) * sizeof(sleep_tree_entry_t)); // FIXME salt
	if (!roots) { free(nodes_vec); errno = ENOMEM; return -1; }
	size_t roots_len = 0;

	for (size_t i = 0; i < nblocks; i++) {
		sleep_tree_entry_t leaf;
		memset(&leaf, 0, sizeof(leaf));
		if (sleep_tree_hash_block(blocks[i], block_lens[i], leaf.hash) != 0) {
			free(nodes_vec);
			free(roots);
			return -1;
		}
		leaf.index = sleep_tree_index(0, (uint64_t)i); // 2 * i ...
		leaf.size = (uint64_t)block_lens[i];
		if (nodes_len >= max_nodes) {
			size_t newcap = max_nodes * 2;
			sleep_tree_entry_t *nvec = (sleep_tree_entry_t *)realloc(nodes_vec, newcap * sizeof(slep_tree_entry_t));
			if (!nvec) { free(nodes_vec); free(roots); errno = ENOMEM; return -1; }
			nodes_vec = nvec;
			max_nodes = newcap;
		}
		nodes_vec[nodes_len++] = leaf;

		if (roots_len >= 64 + 2) {
			size_t newcap = (roots_len + 16);
			sleep_tree_entry_t *r = (sleep_tree_entry_t *)realloc(roots, newcap * sizeof(sleep_tree_entry_t));
			if (!r) { free(nodes_vec); free(roots); errno = ENOMEM; return -1; }
			roots = r;
		}
		roots[roots_len++] = leaf;

		while (roots_len >= 2) {
			sleep_tree_entry_t *a = &roots[roots_len - 2];
			sleep_tree_entry_t *b = &roots[roots_len - 1];

			uint64_t ra = /* FIXME tree depth (a->index) */
			uint64_t rb = /*(b->index)*/
			if (ra != rb) break;
		}
	}
	memcpy(out_nodes, nodes_vec, nodes_len * sizeof(sleep_tree_entry_t));
	*out_count = nodes_len;
	free(nodes_vec);
	free(roots);
	return 0;
}