#include <EEPROM.h>

#include "watchos.h"
#include "Module_Storage.h"

// TODO: Make tiered storage. Smaller allocations can go to RTC RAM.

struct StorageEntry
{
	byte module;
	uint16_t length;
	byte options;
};

#define MODULE_STORAGE_SIZE 3524
#define MODULE_STORAGE_HEADER_SIZE 4
#define MODULE_STORAGE_ENTRY_SIZE 4

#define MODULE_STORAGE_MAGIC 0x6580
#define MODULE_STORAGE_VERSION 0x00

RTC_NOINIT_ATTR byte RTC_MEM[MODULE_STORAGE_SIZE];

byte read_rtc(uint16_t offset);
void write_rtc(uint16_t offset, byte value);

void Module_Storage::initialize()
{
	for (int i = 0; i < MODULE_STORAGE_MAX_ENTRIES; i++)
	{
		m_entry[i] = nullptr;
	}
}

void Module_Storage::start()
{
	load();
}

void Module_Storage::suspend()
{
	commit();
}

bool Module_Storage::isValid()
{
	return ((read_rtc(0) << 8) | read_rtc(1)) == MODULE_STORAGE_MAGIC;
}

void Module_Storage::load()
{
	// Read rtc ram to m_entry
	m_magic = (read_rtc(0) << 8) | read_rtc(1);
	m_version = read_rtc(2);
	m_options = read_rtc(3);

	if (m_magic != MODULE_STORAGE_MAGIC)
		watchos::panic("Invalid magic in persistent data");
	if (m_version != MODULE_STORAGE_VERSION)
		watchos::panic("Invalid version of persistent data");

	for (int i = 0; i < MODULE_STORAGE_MAX_ENTRIES; i++)
	{
		int offset = MODULE_STORAGE_HEADER_SIZE + (MODULE_STORAGE_ENTRY_SIZE * i);
		byte module = read_rtc(offset + 0);
		if (module == 0xFF)
			continue;
		
		m_entry[i] = new StorageEntry();
		m_entry[i]->module = module;
		m_entry[i]->length = (read_rtc(offset + 1) << 8) | read_rtc(offset + 2);
		m_entry[i]->options = read_rtc(offset + 3);
	}
}

void Module_Storage::commit()
{
	// Sync allocations into RTC memory
	write_rtc(0, (byte)(m_magic >> 8));
	write_rtc(1, (byte)(m_magic));
	write_rtc(2, m_version);
	write_rtc(3, m_options);

	for (int i = 0; i < MODULE_STORAGE_MAX_ENTRIES; i++)
	{
		int offset = MODULE_STORAGE_HEADER_SIZE + (MODULE_STORAGE_ENTRY_SIZE * i);
		if (m_entry[i] == nullptr)
		{
			write_rtc(offset + 0, 0xFF);
		}
		else
		{
			write_rtc(offset + 0, m_entry[i]->module);
			write_rtc(offset + 1, (byte)(m_entry[i]->length >> 8));
			write_rtc(offset + 2, (byte)(m_entry[i]->length));
			write_rtc(offset + 3, m_entry[i]->options);
		}
	}
}

void Module_Storage::reset()
{
	m_magic = MODULE_STORAGE_MAGIC;
	m_version = MODULE_STORAGE_VERSION;
	m_options = 0;

	// Write out blank allocations
	write_rtc(0, (byte)(m_magic >> 8));
	write_rtc(1, (byte)(m_magic));
	write_rtc(2, m_version);
	write_rtc(3, m_options);

	for (int i = 0; i < MODULE_STORAGE_MAX_ENTRIES; i++)
	{
		if (m_entry[i] != nullptr)
		{
			free(m_entry[i]);
			m_entry[i] = nullptr;
		}
		
		write_rtc(MODULE_STORAGE_HEADER_SIZE + (MODULE_STORAGE_ENTRY_SIZE * i), 0xFF);
	}
}

bool Module_Storage::allocate(byte module, uint16_t size, byte options)
{
	// Check if an allocation already exists
	for (int i = 0; i < MODULE_STORAGE_MAX_ENTRIES; i++)
	{
		if (m_entry[i] != nullptr && m_entry[i]->module == module)
		{
			if (m_entry[i]->length != size)
				watchos::panic("Storage module does not support reallocating storage with a new size.");
			return false;
		}
	}

	// None exists, allocate
	for (int i = 0; i < MODULE_STORAGE_MAX_ENTRIES; i++)
	{
		if (m_entry[i] == nullptr)
		{
			m_entry[i] = new StorageEntry();
			m_entry[i]->module = module;
			m_entry[i]->length = size;
			m_entry[i]->options = options;
			return true;
		}
	}

	// None exists and we were unable to allocate one
	watchos::panic("Exceeded MODULE_STORAGE_MAX_ENTRIES");
}

void Module_Storage::write(byte module, uint16_t offset, byte value)
{
	// TODO: Check offset doesn't put us past the end of the allocation
	write_rtc(getOffset(module) + offset, value);
}

byte Module_Storage::read(byte module, uint16_t offset)
{
	// TODO: Check offset doesn't put us past the end of the allocation
	return read_rtc(getOffset(module) + offset);
}

bool Module_Storage::allocate(well_known_handle_t module, uint16_t size, byte options)
{
	return allocate(wellKnownHandleToByte(module), size);
}

void Module_Storage::write(well_known_handle_t module, uint16_t offset, byte value)
{
	return write(wellKnownHandleToByte(module), offset, value);
}

byte Module_Storage::read(well_known_handle_t module, uint16_t offset)
{
	return read(wellKnownHandleToByte(module), offset);
}

byte Module_Storage::wellKnownHandleToByte(well_known_handle_t handle)
{
	byte count = 0;
	for (;;)
	{
		if (((handle >> count) & 1) == 1)
			return count;
		else
			count++;
		assert(count < 64);
	}
}

int Module_Storage::getOffset(byte module)
{
	int offset = MODULE_STORAGE_HEADER_SIZE + (MODULE_STORAGE_ENTRY_SIZE * MODULE_STORAGE_MAX_ENTRIES);
	for (int i = 0; i < MODULE_STORAGE_MAX_ENTRIES; i++)
	{
		if (m_entry[i] != nullptr)
		{
			if (m_entry[i]->module == module)
				return offset;
			else
				offset += m_entry[i]->length;
		}
	}
	watchos::panic("Cannot get offset for module %d; not found", module);
}

byte read_rtc(uint16_t offset)
{
	return RTC_MEM[offset];
}

void write_rtc(uint16_t offset, byte value)
{
	RTC_MEM[offset] = value;
}