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

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

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

void mufs_del_tree(struct mufs_tree *tree)
{
	struct mufs *fs;

	if (!tree)
		return;

	fs = tree->fs;

	if (tree == &fs->root)
		return;

	free(tree);
}

static size_t mufs_get_index(struct mufs *fs, uint32_t va, uint8_t depth)
{
	size_t i;

	for (i = 0; i < depth; ++i)
		va >>= fs->log2_nentries;

	va &= BITS(0, fs->log2_nentries);

	return va;
}

static int mufs_do_lookup(struct mufs *fs, uint32_t *page,
	uint32_t base, uint8_t depth, uint32_t va, uint32_t new_page,
	unsigned alloc)
{
	struct flash_dev *dev = fs->dev;
	uint32_t offset, entry;

	if (!base)
		return -1;

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

	offset = mufs_get_index(fs, va, depth) * sizeof entry;
	offset += base << dev->log2_block_size;

	if (flash_read(dev, offset, &entry, sizeof entry) == 0)
		return -1;

	if (!entry) {
		if (alloc) {
			if (depth == 1 && new_page) {
				entry = new_page;
			} else if (mufs_alloc_block(fs, &entry) < 0) {
				return -1;
			}
		}

		if (flash_write(dev, offset, &entry, sizeof entry) == 0)
			return -1;

		if (flash_sync(dev) < 0)
			return -1;
	}

	return mufs_do_lookup(fs, page, entry, depth - 1, va, new_page,
		alloc);
}

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

int mufs_map_page(struct mufs_tree *tree, uint32_t va, uint32_t page)
{
	struct mufs *fs = tree->fs;

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

	return mufs_do_lookup(fs, &page, tree->root, tree->depth, va, page, 1);
}

int mufs_alloc_page(struct mufs_tree *tree, uint32_t *page,
	uint32_t va)
{
	struct mufs *fs = tree->fs;

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

	return mufs_do_lookup(fs, page, tree->root, tree->depth, va, 0, 1);
}

static int mufs_do_free_page(struct mufs *fs, uint32_t *base,
	uint8_t depth, uint32_t va, int free_page)
{
	struct flash_dev *dev = fs->dev;
	uint32_t offset, entry;

	if (!base || *base == 0)
		return -1;

	if (depth == 0) {
		if (free_page) {
			mufs_free_block(fs, *base);
			*base = 0;
		}

		return 0;
	}

	offset = mufs_get_index(fs, va, depth) * sizeof entry;
	offset += (*base << dev->log2_block_size);

	if (flash_read(fs->dev, offset, &entry, sizeof entry) == 0)
		return -1;

	if (mufs_do_free_page(fs, &entry, depth - 1, va, free_page) < 0)
		return -1;

	if (entry != 0)
		return 0;

	if (flash_write0(fs->dev, offset, sizeof entry) == 0)
		return -1;

	if (flash_sync(fs->dev) == 0)
		return -1;

	if (!flash_is_erased(fs->dev, *base, 1))
		return 0;

	if (mufs_free_block(fs, *base) < 0)
		return -1;

	*base = 0;

	return 0;
}

void mufs_unmap_page(struct mufs_tree *tree, uint32_t va)
{
	mufs_do_free_page(tree->fs, &tree->root, tree->depth, va, 0);
}

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

int mufs_extend_tree(struct mufs_tree *tree, uint8_t depth)
{
	struct mufs *fs = tree->fs;
	struct flash_dev *dev = fs->dev;
	uint32_t root;

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

		if (!root)
			continue;

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

		if (flash_sync(dev) < 0) {
			mufs_free_block(fs, root);
			return -1;
		}

		tree->root = root;
	}

	return 0;
}

int mufs_shrink_tree(struct mufs_tree *tree, uint32_t max_size)
{
	struct mufs *fs = tree->fs;
	struct flash_dev *dev = fs->dev;
	size_t index;
	uint32_t base, size;
	uint8_t depth;

	base = tree->root;

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

		if (size < max_size)
			return 0;

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

		if (index <= 1) {
			if (flash_read(dev, (base << dev->log2_block_size), &tree->root,
				sizeof tree->root) == 0)
				return -1;

			--tree->depth;

			mufs_free_block(fs, base);

			continue;
		}

		if (flash_write0(dev, (base << dev->log2_block_size) +
			index * sizeof(uint32_t), (UINT32_C(1) << dev->log2_block_size) -
			index * sizeof(uint32_t)) == 0)
			return -1;

		if (flash_sync(dev) < 0)
			return -1;
	}

	return 0;
}

static size_t mufs_do_read(struct mufs_tree *tree, void *data, uint32_t va,
	size_t len)
{
	struct mufs *fs = tree->fs;
	struct flash_dev *dev = fs->dev;
	uint32_t page, offset;

	offset = va & ((UINT32_C(1) << dev->log2_block_size) - 1);
	va = align(va, dev->log2_block_size);

	memset(data, 0, len);
	len = min(len, (UINT32_C(1) << dev->log2_block_size) - offset);

	if (mufs_lookup_page(tree, &page, va) < 0)
		return 0;

	return flash_read(dev, (page << dev->log2_block_size) + offset, data, len);
}

size_t mufs_tree_read(struct mufs_tree *tree, void *data, uint32_t va, size_t len)
{
	char *buf = data;
	size_t ret, nbytes = 0;

	while (len) {
		if ((ret = mufs_do_read(tree, buf, va, len)) == 0)
			return nbytes;

		buf += ret;
		va += ret;
		nbytes += ret;
		len -= ret;
	}

	return nbytes;
}

static size_t mufs_do_write(struct mufs_tree *tree,
	void *data, uint32_t va, size_t len)
{
	struct mufs *fs = tree->fs;
	struct flash_dev *dev = fs->dev;
	uint32_t page, offset;

	offset = va & ((UINT32_C(1) << dev->log2_block_size) - 1);
	va = align(va, dev->log2_block_size);
	len = min(len, (UINT32_C(1) << dev->log2_block_size) - offset);

	if (mufs_alloc_page(tree, &page, va) < 0)
		return 0;

	return flash_write(dev, (page << dev->log2_block_size) + offset, data,
		len);
}

size_t mufs_tree_write(struct mufs_tree *tree, void *data,
	uint32_t va, size_t len)
{
	uint8_t *buf = data;
	size_t ret, nbytes = 0;

	if (mufs_extend_tree(tree, 1) < 0)
		return 0;

	while (len) {
		if ((ret = mufs_do_write(tree, buf, va, len)) == 0)
			return nbytes;

		buf += ret;
		va += ret;
		nbytes += ret;
		len -= ret;
	}

	return nbytes;
}

int mufs_sync_tree(struct mufs_tree *tree)
{
	struct mufs *fs = tree->fs;
	struct flash_dev *dev = fs->dev;
	struct mufs_dtree dtree;

	dtree.file_size = tree->file_size;
	dtree.root = tree->root;
	dtree.depth = tree->depth;

	if (flash_write(dev, tree->va, &dtree, sizeof dtree) == 0)
		return -1;

	return 0;
}