#include <stdint.h>

#include <bitops.h>
#include <flash.h>

#include <fs/mufs.h>

#include "block.h"
#include "tree.h"

static int mufs_do_lookup(struct mufs *fs, uint32_t *page,
	uint32_t base, uint8_t depth, uint32_t va, unsigned alloc)
{
	char data[1 << fs->dev->log2_block_size];
	uint32_t *table = (uint32_t *)data;
	size_t index;
	uint8_t log2_nentries = fs->dev->log2_block_size - ilog2(sizeof(uint32_t));

	if (!base)
		return -1;

	if (!depth) {
		*page = base;
		return 0;
	}

	if (flash_read(fs->dev, base << fs->dev->log2_block_size, data,
		sizeof data) == 0)
		return -1;

	index = va & ((1 << (log2_nentries * (depth - 1))) - 1);

	if (!table[index]) {
		if (!alloc || mufs_alloc_block(fs, &table[index]) < 0)
			return -1;
	}

	return mufs_do_lookup(fs, page, table[index], depth - 1, va, alloc);
}

int mufs_lookup_page(struct mufs *fs, struct mufs_tree *tree, uint32_t *page,
	uint32_t va)
{
	return mufs_do_lookup(fs, page, tree->root, tree->depth, va, 0);
}

int mufs_alloc_page(struct mufs *fs, struct mufs_tree *tree, uint32_t *page,
	uint32_t va)
{
	return mufs_do_lookup(fs, page, tree->root, tree->depth, va, 1);
}

static int mufs_do_free_page(struct mufs *fs, uint32_t base,
	uint8_t depth, uint32_t va)
{
	char data[1 << fs->dev->log2_block_size];
	uint32_t *table = (uint32_t *)data;
	size_t index;
	uint8_t log2_nentries = fs->dev->log2_block_size - ilog2(sizeof(uint32_t));

	if (!base)
		return -1;

	if (!depth) {
		mufs_free_block(fs, base);
		return 0;
	}

	if (flash_read(fs->dev, base << fs->dev->log2_block_size, data,
		sizeof data) == 0)
		return -1;

	index = va & ((1 << (log2_nentries * (depth - 1))) - 1);

	if (mufs_do_free_page(fs, table[index], depth - 1, va) < 0)
		return -1;

	table[index] = 0;

	for (index = 0; index < (UINT32_C(1) << log2_nentries); ++index) {
		if (!table[index])
			continue;

		if (flash_write(fs->dev, base << fs->dev->log2_block_size, data,
			sizeof data) == 0)
			return -1;

		return 0;
	}

	return mufs_free_block(fs, base);
}

void mufs_free_page(struct mufs *fs, struct mufs_tree *tree, uint32_t va)
{
	mufs_do_free_page(fs, tree->root, tree->depth, va);
}

int mufs_extend_tree(struct mufs *fs, struct mufs_tree *tree, uint8_t depth)
{
	uint32_t root;

	for (; tree->depth <= depth; ++tree->depth) {
		if (mufs_alloc_block(fs, &root) < 0)
			return -1;

		if (!root)
			continue;

		if (flash_write(fs->dev, root << fs->dev->log2_block_size, &tree->root,
			sizeof tree->root) == 0) {
			mufs_free_block(fs, root);
			return -1;
		}
	}

	return 0;
}

int mufs_shrink_tree(struct mufs *fs, struct mufs_tree *tree, uint32_t max_size)
{
	char data[1 << fs->dev->log2_block_size];
	uint32_t *table = (uint32_t *)data;
	size_t index;
	uint32_t base, size;
	uint8_t log2_nentries = fs->dev->log2_block_size - ilog2(sizeof(uint32_t));
	uint8_t depth;

	base = tree->root;

	for (depth = tree->depth; depth; --depth) {
		size = 1 << (depth * log2_nentries + fs->dev->log2_block_size);

		if (size < max_size)
			return 0;

		index = max_size >> ((depth - 1) * log2_nentries +
			fs->dev->log2_block_size);

		if (flash_read(fs->dev, base, data, sizeof data) == 0)
			return -1;

		if (index <= 1) {
			tree->root = table[0];
			--tree->depth;

			mufs_free_block(fs, base);

			continue;
		}

		for (; index < (UINT32_C(1) << log2_nentries); ++index) {
			table[index] = 0;
		}

		if (flash_write(fs->dev, base << fs->dev->log2_block_size, data,
			sizeof data) == 0)
			return -1;
	}

	return 0;
}