/*
 * Copyright (c) 2016-2017, ARM Limited and Contributors. All rights reserved.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

#include <assert.h>
#include <debug.h>
#include <errno.h>
#include <io_block.h>
#include <io_driver.h>
#include <io_storage.h>
#include <platform_def.h>
#include <string.h>
#include <utils.h>

typedef struct {
	io_block_dev_spec_t	*dev_spec;
	uintptr_t		base;
	size_t			file_pos;
	size_t			size;
} block_dev_state_t;

#define is_power_of_2(x)	((x != 0) && ((x & (x - 1)) == 0))

io_type_t device_type_block(void);

static int block_open(io_dev_info_t *dev_info, const uintptr_t spec,
		      io_entity_t *entity);
static int block_seek(io_entity_t *entity, int mode, ssize_t offset);
static int block_read(io_entity_t *entity, uintptr_t buffer, size_t length,
		      size_t *length_read);
static int block_write(io_entity_t *entity, const uintptr_t buffer,
		       size_t length, size_t *length_written);
static int block_close(io_entity_t *entity);
static int block_dev_open(const uintptr_t dev_spec, io_dev_info_t **dev_info);
static int block_dev_close(io_dev_info_t *dev_info);

static const io_dev_connector_t block_dev_connector = {
	.dev_open	= block_dev_open
};

static const io_dev_funcs_t block_dev_funcs = {
	.type		= device_type_block,
	.open		= block_open,
	.seek		= block_seek,
	.size		= NULL,
	.read		= block_read,
	.write		= block_write,
	.close		= block_close,
	.dev_init	= NULL,
	.dev_close	= block_dev_close,
};

static block_dev_state_t state_pool[MAX_IO_BLOCK_DEVICES];
static io_dev_info_t dev_info_pool[MAX_IO_BLOCK_DEVICES];

/* Track number of allocated block state */
static unsigned int block_dev_count;

io_type_t device_type_block(void)
{
	return IO_TYPE_BLOCK;
}

/* Locate a block state in the pool, specified by address */
static int find_first_block_state(const io_block_dev_spec_t *dev_spec,
				  unsigned int *index_out)
{
	int result = -ENOENT;
	for (int index = 0; index < MAX_IO_BLOCK_DEVICES; ++index) {
		/* dev_spec is used as identifier since it's unique */
		if (state_pool[index].dev_spec == dev_spec) {
			result = 0;
			*index_out = index;
			break;
		}
	}
	return result;
}

/* Allocate a device info from the pool and return a pointer to it */
static int allocate_dev_info(io_dev_info_t **dev_info)
{
	int result = -ENOMEM;
	assert(dev_info != NULL);

	if (block_dev_count < MAX_IO_BLOCK_DEVICES) {
		unsigned int index = 0;
		result = find_first_block_state(NULL, &index);
		assert(result == 0);
		/* initialize dev_info */
		dev_info_pool[index].funcs = &block_dev_funcs;
		dev_info_pool[index].info = (uintptr_t)&state_pool[index];
		*dev_info = &dev_info_pool[index];
		++block_dev_count;
	}

	return result;
}


/* Release a device info to the pool */
static int free_dev_info(io_dev_info_t *dev_info)
{
	int result;
	unsigned int index = 0;
	block_dev_state_t *state;
	assert(dev_info != NULL);

	state = (block_dev_state_t *)dev_info->info;
	result = find_first_block_state(state->dev_spec, &index);
	if (result ==  0) {
		/* free if device info is valid */
		zeromem(state, sizeof(block_dev_state_t));
		zeromem(dev_info, sizeof(io_dev_info_t));
		--block_dev_count;
	}

	return result;
}

static int block_open(io_dev_info_t *dev_info, const uintptr_t spec,
		      io_entity_t *entity)
{
	block_dev_state_t *cur;
	io_block_spec_t *region;

	assert((dev_info->info != (uintptr_t)NULL) &&
	       (spec != (uintptr_t)NULL) &&
	       (entity->info == (uintptr_t)NULL));

	region = (io_block_spec_t *)spec;
	cur = (block_dev_state_t *)dev_info->info;
	assert(((region->offset % cur->dev_spec->block_size) == 0) &&
	       ((region->length % cur->dev_spec->block_size) == 0));

	cur->base = region->offset;
	cur->size = region->length;
	cur->file_pos = 0;

	entity->info = (uintptr_t)cur;
	return 0;
}

/* parameter offset is relative address at here */
static int block_seek(io_entity_t *entity, int mode, ssize_t offset)
{
	block_dev_state_t *cur;

	assert(entity->info != (uintptr_t)NULL);

	cur = (block_dev_state_t *)entity->info;
	assert((offset >= 0) && (offset < cur->size));

	switch (mode) {
	case IO_SEEK_SET:
		cur->file_pos = offset;
		break;
	case IO_SEEK_CUR:
		cur->file_pos += offset;
		break;
	default:
		return -EINVAL;
	}
	assert(cur->file_pos < cur->size);
	return 0;
}

static int block_read(io_entity_t *entity, uintptr_t buffer, size_t length,
		      size_t *length_read)
{
	block_dev_state_t *cur;
	io_block_spec_t *buf;
	io_block_ops_t *ops;
	size_t aligned_length, skip, count, left, padding, block_size;
	int lba;
	int buffer_not_aligned;

	assert(entity->info != (uintptr_t)NULL);
	cur = (block_dev_state_t *)entity->info;
	ops = &(cur->dev_spec->ops);
	buf = &(cur->dev_spec->buffer);
	block_size = cur->dev_spec->block_size;
	assert((length <= cur->size) &&
	       (length > 0) &&
	       (ops->read != 0));

	if ((buffer & (block_size - 1)) != 0) {
		/*
		 * buffer isn't aligned with block size.
		 * Block device always relies on DMA operation.
		 * It's better to make the buffer as block size aligned.
		 */
		buffer_not_aligned = 1;
	} else {
		buffer_not_aligned = 0;
	}

	skip = cur->file_pos % block_size;
	aligned_length = ((skip + length) + (block_size - 1)) &
			 ~(block_size - 1);
	padding = aligned_length - (skip + length);
	left = aligned_length;
	do {
		lba = (cur->file_pos + cur->base) / block_size;
		if (left >= buf->length) {
			/*
			 * Since left is larger, it's impossible to padding.
			 *
			 * If buffer isn't aligned, we need to use aligned
			 * buffer instead.
			 */
			if (skip || buffer_not_aligned) {
				/*
				 * The beginning address (file_pos) isn't
				 * aligned with block size, we need to use
				 * block buffer to read block. Since block
				 * device is always relied on DMA operation.
				 */
				count = ops->read(lba, buf->offset,
						  buf->length);
			} else {
				count = ops->read(lba, buffer, buf->length);
			}
			assert(count == buf->length);
			cur->file_pos += count - skip;
			if (skip || buffer_not_aligned) {
				/*
				 * Since there's not aligned block size caused
				 * by skip or not aligned buffer, block buffer
				 * is used to store data.
				 */
				memcpy((void *)buffer,
				       (void *)(buf->offset + skip),
				       count - skip);
			}
			left = left - (count - skip);
		} else {
			if (skip || padding || buffer_not_aligned) {
				/*
				 * The beginning address (file_pos) isn't
				 * aligned with block size, we have to read
				 * full block by block buffer instead.
				 * The size isn't aligned with block size.
				 * Use block buffer to avoid overflow.
				 *
				 * If buffer isn't aligned, use block buffer
				 * to avoid DMA error.
				 */
				count = ops->read(lba, buf->offset, left);
			} else
				count = ops->read(lba, buffer, left);
			assert(count == left);
			left = left - (skip + padding);
			cur->file_pos += left;
			if (skip || padding || buffer_not_aligned) {
				/*
				 * Since there's not aligned block size or
				 * buffer, block buffer is used to store data.
				 */
				memcpy((void *)buffer,
				       (void *)(buf->offset + skip),
				       left);
			}
			/* It's already the last block operation */
			left = 0;
		}
		skip = cur->file_pos % block_size;
	} while (left > 0);
	*length_read = length;

	return 0;
}

static int block_write(io_entity_t *entity, const uintptr_t buffer,
		       size_t length, size_t *length_written)
{
	block_dev_state_t *cur;
	io_block_spec_t *buf;
	io_block_ops_t *ops;
	size_t aligned_length, skip, count, left, padding, block_size;
	int lba;
	int buffer_not_aligned;

	assert(entity->info != (uintptr_t)NULL);
	cur = (block_dev_state_t *)entity->info;
	ops = &(cur->dev_spec->ops);
	buf = &(cur->dev_spec->buffer);
	block_size = cur->dev_spec->block_size;
	assert((length <= cur->size) &&
	       (length > 0) &&
	       (ops->read != 0) &&
	       (ops->write != 0));

	if ((buffer & (block_size - 1)) != 0) {
		/*
		 * buffer isn't aligned with block size.
		 * Block device always relies on DMA operation.
		 * It's better to make the buffer as block size aligned.
		 */
		buffer_not_aligned = 1;
	} else {
		buffer_not_aligned = 0;
	}

	skip = cur->file_pos % block_size;
	aligned_length = ((skip + length) + (block_size - 1)) &
			 ~(block_size - 1);
	padding = aligned_length - (skip + length);
	left = aligned_length;
	do {
		lba = (cur->file_pos + cur->base) / block_size;
		if (left >= buf->length) {
			/* Since left is larger, it's impossible to padding. */
			if (skip || buffer_not_aligned) {
				/*
				 * The beginning address (file_pos) isn't
				 * aligned with block size or buffer isn't
				 * aligned, we need to use block buffer to
				 * write block.
				 */
				count = ops->read(lba, buf->offset,
						  buf->length);
				assert(count == buf->length);
				memcpy((void *)(buf->offset + skip),
				       (void *)buffer,
				       count - skip);
				count = ops->write(lba, buf->offset,
						   buf->length);
			} else
				count = ops->write(lba, buffer, buf->length);
			assert(count == buf->length);
			cur->file_pos += count - skip;
			left = left - (count - skip);
		} else {
			if (skip || padding || buffer_not_aligned) {
				/*
				 * The beginning address (file_pos) isn't
				 * aligned with block size, we need to avoid
				 * poluate data in the beginning. Reading and
				 * skipping the beginning is the only way.
				 * The size isn't aligned with block size.
				 * Use block buffer to avoid overflow.
				 *
				 * If buffer isn't aligned, use block buffer
				 * to avoid DMA error.
				 */
				count = ops->read(lba, buf->offset, left);
				assert(count == left);
				memcpy((void *)(buf->offset + skip),
				       (void *)buffer,
				       left - skip - padding);
				count = ops->write(lba, buf->offset, left);
			} else
				count = ops->write(lba, buffer, left);
			assert(count == left);
			cur->file_pos += left - (skip + padding);
			/* It's already the last block operation */
			left = 0;
		}
		skip = cur->file_pos % block_size;
	} while (left > 0);
	*length_written = length;
	return 0;
}

static int block_close(io_entity_t *entity)
{
	entity->info = (uintptr_t)NULL;
	return 0;
}

static int block_dev_open(const uintptr_t dev_spec, io_dev_info_t **dev_info)
{
	block_dev_state_t *cur;
	io_block_spec_t *buffer;
	io_dev_info_t *info;
	size_t block_size;
	int result;

	assert(dev_info != NULL);
	result = allocate_dev_info(&info);
	if (result)
		return -ENOENT;

	cur = (block_dev_state_t *)info->info;
	/* dev_spec is type of io_block_dev_spec_t. */
	cur->dev_spec = (io_block_dev_spec_t *)dev_spec;
	buffer = &(cur->dev_spec->buffer);
	block_size = cur->dev_spec->block_size;
	assert((block_size > 0) &&
	       (is_power_of_2(block_size) != 0) &&
	       ((buffer->offset % block_size) == 0) &&
	       ((buffer->length % block_size) == 0));

	*dev_info = info;	/* cast away const */
	(void)block_size;
	(void)buffer;
	return 0;
}

static int block_dev_close(io_dev_info_t *dev_info)
{
	return free_dev_info(dev_info);
}

/* Exported functions */

/* Register the Block driver with the IO abstraction */
int register_io_dev_block(const io_dev_connector_t **dev_con)
{
	int result;

	assert(dev_con != NULL);

	/*
	 * Since dev_info isn't really used in io_register_device, always
	 * use the same device info at here instead.
	 */
	result = io_register_device(&dev_info_pool[0]);
	if (result == 0)
		*dev_con = &block_dev_connector;
	return result;
}