#include <ZSimulation/ZPropertyBuffer.hpp>
#include <ZSimulation/ZNetworkUpdateStream.hpp>

#include <SST/SST_SysMem.h>

#define BUFFER_SYNC	 0
#define BUFFER_LOCAL 1
#define BUFFER_WRITE 0
#define BUFFER_READ  1

#define MIN_ALLOC	 4 

struct PropertyPage {
	eID				Id;			// id of the entity that these properties are allocated for
	void*			Data[2];	// the data for this page (read and write)
	size_t			Size;		// amount of data on this page
	size_t			Offset;		// the current offset on this page
	PropertyPage*	Next;		// the next page, or NULL

	ZArray<ZPair<size_t, size_t>>						Avail;		// available data slots on this page
	ZArray<ZPair<ZPropertyBuffer::PropertyKey, size_t>> Updates;	// properties that have been updated since last update

	// c'tor
	PropertyPage(size_t alloc_size) : Next(NULL) { 
		Data[0] = SST_OS_AllocPages(alloc_size*2);		// allocate 2x to account for read and write
		Data[1] = (char*)Data[0] + alloc_size;			// split in half, with first half being write and second for read
		Size	= alloc_size;
		Offset	= 0; 
	}
	
	// d'tor
	~PropertyPage() { 
		delete Next; 
		if (Size > 0) {
			SST_OS_FreePages(Data[0], Size*2); 
		}
	}
};

/*************************************************************************/

// Key is currently a 64 bit integer, which is partitioned as follows
//
// [ SYNC MOE (2) | OFFSET (14) | SUB PAGE (16) | PAGE (32) ]

static ZPropertyBuffer::PropertyKey GenerateKey(ZPropertyBuffer::SyncMode mode, uint32_t page, uint16_t sub_page, uint16_t offset) {

	uint64_t key;
	
	key = mode;
	key = key << 14;
	key = key + offset;
	key = key << 16;
	key = key + sub_page;
	key = key << 32;
	key = key + page;

	// dumb assert
	static_assert(sizeof(eID) == 4, "Cannot generate property key with sizeof(eID) not 32 bits, update property buffer");

	return key;
}

static void ParseKey(ZPropertyBuffer::PropertyKey key, ZPropertyBuffer::SyncMode* mode, uint32_t* page, uint16_t* sub_page, uint16_t* offset) {
	*page	  = (uint32_t)(key & 0xFFFFFFFF);
	*sub_page = (uint16_t)((key >> 32) & 0xFFFF);
	*offset	  = (uint16_t)((key >> 48) & 0x3FFF);
	*mode	  = (ZPropertyBuffer::SyncMode)((key >> 62) & 0xFFFF);
}

void* ZPropertyBuffer::GetData( PropertyKey key, int rw, PropertyPage** page_out /*= NULL*/, SyncMode* mode_out /*= NULL*/ )
{
	SyncMode mode;
	uint32_t page_idx;
	uint16_t sub_page_idx, offset;

	ParseKey(key, &mode, &page_idx, &sub_page_idx, &offset);				// parse data key
	
	ZArray<PropertyPage*>& pages = mode != SYNC_NONE ? Pages[BUFFER_SYNC] : 
													   Pages[BUFFER_LOCAL];	// determine if we are using sync or local data
	
	if (pages.Size() >= page_idx) {											// determine if we have this many pages
		PropertyPage* page = pages[page_idx];

		for (size_t i = 0; i < sub_page_idx; i++) {							// see of we can find the sub page
			page = page->Next;
			if (page == NULL) {												// check for bad sub page index
				return NULL;
			}
		}

		if (offset % MIN_ALLOC == 0 && offset <= page->Size - MIN_ALLOC) {	// see if our offset is valid
			if (page_out != NULL) *page_out = page;							// get pointer to page
			if (mode_out != NULL) *mode_out = mode;							// get sync value
			return (char*)page->Data[rw] + offset;							// get read / write pointer
		} else return NULL;
	} else return NULL;
}

/*************************************************************************/

ZPropertyBuffer::Property ZPropertyBuffer::AllocData( PropertyPage* page, eID id, uint32_t page_idx, size_t size, SyncMode mode)
{
	Property prop;

	ZArray<PropertyPage*>& pages = mode ? Pages[BUFFER_SYNC] : Pages[BUFFER_LOCAL];			// determine if we are using sync or local data

	size_t alloc_size = size > MIN_ALLOC ? (size + MIN_ALLOC - 1) & ~(MIN_ALLOC - 1)		// ensure we allocate at least MIN_ALLOC, and align to MIN_ALLOC
										 :  size;								

	uint16_t sub_page_idx;

	for (sub_page_idx = 0; page != NULL; sub_page_idx++) {									// loop pages seeing if we can drop the data in here
		if (page->Offset + size <= page->Size) {											// see if we have room
			uint16_t offset = (uint16_t)page->Offset;										// we do, so get current offset
			page->Offset	= page->Offset + alloc_size;									// increment stored offset by alloc size, not property size
			prop.Key		= GenerateKey(mode, page_idx, sub_page_idx, offset);			// key for this value
			prop.Write		= (char*)page->Data[BUFFER_WRITE] + offset;						// write pointer for this value
			prop.Read		= (char*)page->Data[BUFFER_READ] + offset;						// read pointer for this value
			prop.Size		= size;															// size of this value
			return prop;
		} else {
			for (size_t i = 0; i < page->Avail.Size(); i++) {								// for each available slot from previous deallocs
				if (page->Avail[i].Second >= size) {										// see if we have room
					uint16_t offset = (uint16_t)page->Avail[i].First;						// we do, so store the offset
					page->Avail.Erase(i);													// no longer available
					prop.Key	    = GenerateKey(mode, page_idx, sub_page_idx, offset);	// key for this value
					prop.Write	    = (char*)page->Data[BUFFER_WRITE] + offset;				// write pointer for this value
					prop.Read		= (char*)page->Data[BUFFER_READ] + offset;				// read pointer for this value
					prop.Size		= size;													// size of this value
					return prop;
				}
			}

			if (page->Next != NULL) {
				page = page->Next;															// NEXT
			} else {
				sub_page_idx++;																// sub page index will be next index
				break;
			}
		}
	}

	// found nothing, time to allocate new

	if (Pages[BUFFER_LOCAL].Size() <= page_idx) {					// push NULL entry for local buffer if not present at page_idx
		Pages[BUFFER_LOCAL].PushBack(NULL);
	}

	if (Pages[BUFFER_SYNC].Size() <= page_idx) {					// push NULL entry for sync buffer if not present at page_idx
		Pages[BUFFER_SYNC].PushBack(NULL);
	}

	const uint32_t page_size		= SST_OS_GetPageSize();			// system page size
	const uint32_t min_buffer_alloc = page_size / 2;				// minimum amount we will allocate
	
	uint32_t page_alloc_size		= min_buffer_alloc;				// starting page alloc
	
	if (size > page_alloc_size) {
		page_alloc_size = 
			((size+min_buffer_alloc-1) & ~(min_buffer_alloc-1));	// round up to nearest multiple of min_buffer_alloc
	}

	PropertyPage* npage = new PropertyPage(page_alloc_size);		// new page
	npage->Id			= id;										// set id for page
	npage->Offset		= alloc_size;								// allocate enough space for the first property

	if (page != NULL) {												// link page into page list
		page->Next = npage;
	} else {														// allocating new pages
		pages[page_idx] = npage;
	}
	
	prop.Key   = GenerateKey(mode, page_idx, sub_page_idx, 0);		// generate a key for the correct page with sub-page and offset 0
	prop.Write = (char*)npage->Data[BUFFER_WRITE];					// set the write pointer
	prop.Read  = (char*)npage->Data[BUFFER_READ];					// set the read pointer
	prop.Size  = size;												// set the size of the data
	return prop;
}

/*************************************************************************/

ZPropertyBuffer::ZPropertyBuffer()
{

}

/*************************************************************************/

ZPropertyBuffer::~ZPropertyBuffer()
{
	for (size_t i = 0; i < Pages[BUFFER_LOCAL].Size(); i++) {
		delete Pages[BUFFER_LOCAL][i];
	}

	for (size_t i = 0; i < Pages[BUFFER_SYNC].Size(); i++) {
		delete Pages[BUFFER_SYNC][i];
	}
}

/*************************************************************************/

ZPropertyBuffer::Property ZPropertyBuffer::AllocProperty(eID id, size_t size, SyncMode mode)
{
	// sanity check - synchronized properties should be less than ZSIM_BUFFER_SYNC_MAX
	ZASSERT_RUNTIME(!mode || size < ZSIM_SYNCBUFFER_MAX, "Synchronized property allocation too large!");

	auto itr = IdPageMap.Find(id);

	ZArray<PropertyPage*>& pages = mode != SYNC_NONE ? Pages[BUFFER_SYNC] : 
													   Pages[BUFFER_LOCAL];	

	if (itr != IdPageMap.End()) {
		uint32_t page_idx  = itr.GetValue();
		PropertyPage* page = pages[page_idx];

		return AllocData(page, id, page_idx, size, mode);
	}

	if (Avail.Size() > 0) {
		size_t page_idx = Avail.PopBack();
		IdPageMap.Put(id, page_idx);
		return AllocData(NULL, id, page_idx, size, mode);
	} else {
		IdPageMap.Put(id, pages.Size());
		return AllocData(NULL, id, pages.Size(), size, mode);
	}
}

/*************************************************************************/

void ZPropertyBuffer::DeallocEntity( eID id )
{
	auto itr = IdPageMap.Find(id);

	if (itr != IdPageMap.End()) {
		uint32_t page_idx  = itr.GetValue();

		delete Pages[BUFFER_LOCAL][page_idx];
		delete Pages[BUFFER_SYNC][page_idx];

		Pages[BUFFER_LOCAL][page_idx] = NULL;
		Pages[BUFFER_SYNC][page_idx] = NULL;

		Avail.PushBack(page_idx);

		IdPageMap.Erase(id);
	}
}

/*************************************************************************/

void ZPropertyBuffer::DeallocProperty( PropertyKey key, size_t size )
{
	SyncMode mode;
	uint32_t page_idx;
	uint16_t sub_page_idx, offset;

	ParseKey(key, &mode, &page_idx, &sub_page_idx, &offset);
		
	ZArray<PropertyPage*>& pages = mode != SYNC_NONE ? Pages[BUFFER_SYNC] : 
													   Pages[BUFFER_LOCAL];	

	PropertyPage* page = pages[page_idx];

	for (size_t i = 0; i < sub_page_idx; i++) {
		page = page->Next;
	}

	page->Avail.PushBack(ZPair<size_t, size_t>(offset, size));
}

/*************************************************************************/

void ZPropertyBuffer::FlagModified( PropertyKey key, size_t size )
{
	PropertyPage* page;
	SyncMode	  mode;

	void* data = GetData(key, BUFFER_WRITE, &page, &mode);

	ZASSERT(data != NULL && page != NULL, "Attempt to flag modified on invalid property key");

	page->Updates.PushBack(ZPair<PropertyKey, size_t>(key, size));
}

/*************************************************************************/

void ZPropertyBuffer::UpdateModified(bool sync, ZNetworkUpdateStream& stream)
{
	// wherever we have an update, copy the write buffer
	for (size_t b = 0; b < 2; b++) {
		for (size_t i = 0; i < Pages[b].Size(); i++) {
			PropertyPage* page = Pages[b][i];
			ZPropertyBufferUpdate* buffer_update = NULL;

			// synchronized buffers require a property buffer update call
			if (sync && b == BUFFER_SYNC && page != NULL && page->Updates.Size() > 0) {
				buffer_update = stream.AllocUpdate(page->Id);
			}

			// for all pages and sub-pages
			while (page != NULL) {
				for (size_t j = 0; page != NULL && j < page->Updates.Size(); j++) {
					PropertyKey key = page->Updates[j].First;
					size_t size		= page->Updates[j].Second; 

					SyncMode mode;
					uint32_t page_idx;
					uint16_t sub_page_idx, offset;

					ParseKey(key, &mode, &page_idx, &sub_page_idx, &offset);

					void* data = (char*)page->Data[BUFFER_WRITE] + offset;
					MemCopyChar((char*)page->Data[BUFFER_READ] + offset, data, size);

					// if needed, make it happen
					if (buffer_update != NULL && mode == SYNC_SEND) {
						ZPropertyBufferUpdate::Update update;

						update.SubPage = sub_page_idx;
						update.Offset  = offset;
						update.Size	   = size;
						update.Data	   = (char*)page->Data[BUFFER_READ] + offset;

						buffer_update->Updates.PushBack(update);
					}
				}

				page->Updates.Clear();
				page = page->Next;
			}

			// push the entity update
			if (buffer_update != NULL) {
				stream.EnqueueOutgoingUpdate(buffer_update);
			}
		}
	}
}

/*************************************************************************/

void ZPropertyBuffer::UpdateModified( eID id, bool sync, ZNetworkUpdateStream& stream )
{
	auto itr = IdPageMap.Find(id);
	
	if (itr != IdPageMap.End()) {
		uint32_t page_idx = itr.GetValue();

		// wherever we have an update, copy the write buffer
		for (size_t b = BUFFER_LOCAL; b != BUFFER_SYNC; b = BUFFER_SYNC) {
			PropertyPage* page = Pages[b][page_idx];
			ZPropertyBufferUpdate* buffer_update = NULL;

			// synchronized buffers require a property buffer update call
			if (sync && b == BUFFER_SYNC && page != NULL) {
				buffer_update = stream.AllocUpdate(id);
			}

			// for all pages and sub-pages
			while (page != NULL) {
				for (size_t j = 0; page != NULL && j < page->Updates.Size(); j++) {
					PropertyKey key = page->Updates[j].First;
					size_t size		= page->Updates[j].Second; 

					SyncMode mode;
					uint32_t page_idx;
					uint16_t sub_page_idx, offset;

					ParseKey(key, &mode, &page_idx, &sub_page_idx, &offset);

					void* data = (char*)page->Data[BUFFER_WRITE] + offset;
					MemCopyChar((char*)page->Data[BUFFER_READ] + offset, data, size);

					if (buffer_update != NULL && mode == SYNC_SEND) {
						ZPropertyBufferUpdate::Update update;

						update.SubPage = sub_page_idx;
						update.Offset  = offset;
						update.Size	   = size;
						update.Data	   = (char*)page->Data[BUFFER_READ] + offset;

						buffer_update->Updates.PushBack(update);
					}
				}

				page->Updates.Clear();
				page = page->Next;
			}

			// push the entity update
			if (buffer_update != NULL) {
				stream.EnqueueOutgoingUpdate(buffer_update);
			}
		}
	} else {
		SystemLogError("Unable to update entity property buffer - invalid id");
	}
}

/*************************************************************************/

static void* GetSyncData(PropertyPage* page, uint16_t offset)
{
	if (offset % MIN_ALLOC == 0 && offset <= page->Size - MIN_ALLOC) {	// see if our offset is valid
		return (char*)page->Data[BUFFER_READ] + offset;					// return read
	} else return NULL;
}

void ZPropertyBuffer::ApplyBufferUpdates( const ZArray<ZPropertyBufferUpdate*>& updates )
{
	for (size_t i = 0; i < updates.Size(); i++) {
		ZPropertyBufferUpdate& buffer_update = *updates[i];
		eID id = buffer_update.Id;

		ZArray<PropertyPage*>& pages = Pages[BUFFER_SYNC];						// always sync data

		auto itr = IdPageMap.Find(id);											// find page mapped for id

		if (itr != IdPageMap.End()) {											// make sure we have a valid one
			uint32_t	  page_idx  = itr.GetValue();			
			PropertyPage* base_page = pages[page_idx];

			for (size_t j = 0; j < buffer_update.Updates.Size(); j++) {
				ZPropertyBufferUpdate::Update& update = buffer_update.Updates[j];

				uint16_t sub_page_idx  = update.SubPage;
				uint16_t offset		   = update.Offset;
				size_t	 size		   = update.Size;

				PropertyPage* sub_page = base_page;

				for (uint16_t k = 0; k < sub_page_idx; k++) {
					if (sub_page->Next != NULL) {
						sub_page = sub_page->Next;
					}
				}

				if (sub_page != NULL) {
					void* rd_ptr = GetSyncData(sub_page, offset);

					if (rd_ptr != NULL) {
						MemCopyChar(rd_ptr, update.Data, size);
					} else {
						SystemLogError("Recieved invalid buffer update offset");
					}
				} else {
					SystemLogError("Recieved invalid buffer update sub page");
				}
			}
		}
	}
}