#include "../../Core/DataStore/RuDataStore_Disk.h"

#pragma unmanaged

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

ruRTTI_DEFINE_SUBCLASS(CRuDataStore_Disk, IRuObject)
ruHEAP_DEFINE_SUBCLASS(CRuDataStore_Disk, "CRuDataStore_Disk", "CRuDataStore_Base")
ruCLASSGUID_DEFINE(CRuDataStore_Disk, 0x319CB9B4, 0x4CB94600, 0xA36D2022, 0xD059DF1D)

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

BOOL CRuDataStore_Disk::KeyStruct::operator==(const KeyStruct &key) const
{
	return (memcmp(m_key, key.m_key, m_keyLength) == 0);
}

BOOL CRuDataStore_Disk::KeyStruct::operator<(const KeyStruct &key) const
{
	return (memcmp(m_key, key.m_key, m_keyLength) < 0);
}

BOOL CRuDataStore_Disk::KeyStruct::operator>(const KeyStruct &key) const
{
	return (memcmp(m_key, key.m_key, m_keyLength) > 0);
}

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

CRuDataStore_Disk::CRuDataStore_Disk()
:	m_ioBufferSize(4096),
	m_tableListPageSize(32),
	m_tablePageSize(64),
	m_stream(NULL),
	m_ioBuffer(NULL)
{
	m_ioBuffer = ruNEW BYTE [m_ioBufferSize];
}

CRuDataStore_Disk::~CRuDataStore_Disk()
{
	// Flush any changes to disk
	Flush();

	for(INT32 i = 0; i < m_tableHeaders.Count(); ++i)
	{
		// Iterate over all stream headers and delete each
		m_tableHeaders[i]->m_streamHeaders.ForAllElements(boost::bind(&CRuDataStore_Disk::DeleteStreamHeadersCallback, this, _1, _2));

		// Clear binary tree
		m_tableHeaders[i]->m_streamHeaders.Clear();

		// Delete the table header
		delete m_tableHeaders[i];
	}

	// Free resources
	delete m_stream;
	delete[] m_ioBuffer;
}

BOOL CRuDataStore_Disk::DeleteStreamHeadersCallback(const void *key, void *data)
{
	BYTE *headerKey = *((BYTE **) key);

	delete[] headerKey;

	return TRUE;
}

BOOL CRuDataStore_Disk::Create(const char *dataStoreName)
{
	m_readOnly = FALSE;

	// Create data store stream
	m_stream = ruNEW CRuFileStream();
	m_stream->Create(dataStoreName);

	m_dataStoreHeader.m_tableListPosition = 0;
	m_dataStoreHeader.m_tableListSize = 0;
	m_dataStoreHeader.m_tableListAllocationSize = 0;

	// Reserve space for data store header
	SaveDataStoreHeader(&m_dataStoreHeader);

	// Create initial table list
	SaveTableHeaders(m_dataStoreHeader, m_tableHeaders);

	// Update data store header
	SaveDataStoreHeader(&m_dataStoreHeader);

	return TRUE;
}

BOOL CRuDataStore_Disk::Open(const char *dataStoreName, BOOL readOnly)
{
	m_readOnly = readOnly;

	RuStreamOpenMethod openMethod = ruSOM_OpenReadWrite;
	if(m_readOnly)
	{
		openMethod = ruSOM_OpenRead;
	}

	// Open data store stream
	m_stream = ruNEW CRuFileStream();

	if(m_stream->Open(dataStoreName, openMethod) == ruSEC_OK)
	{
		// Load data store header
		LoadDataStoreHeader(&m_dataStoreHeader);

		// Load table list
		LoadTableHeaders(m_dataStoreHeader, m_tableHeaders);

		// Load stream lists
		for(INT32 i = 0; i < m_tableHeaders.Count(); ++i)
		{
			LoadStreamHeaders(m_tableHeaders, i);
		}

		return TRUE;
	}

	delete m_stream;
	m_stream = NULL;

	return FALSE;
}

INT32 CRuDataStore_Disk::CreateTable(const char *tableName, INT32 keyLength)
{
	if(m_readOnly)
	{
		return -2;
	}

	// Check whether the specified table name is already in the list
	for(INT32 i = 0; i < m_tableHeaders.Count(); ++i)
	{
		if(strcmpi(tableName, m_tableHeaders[i]->m_name) == 0)
		{
			// If the table already exists with the correct key length, return its handle
			if(m_tableHeaders[i]->m_keyLength == keyLength)
			{
				return i;
			}

			// Otherwise, return an error
			return -2;
		}
	}

	// If the table is to be created, append it to the table list
	INT32 tableHandle = m_tableHeaders.Count();
	m_tableHeaders.Add(ruNEW TableHeader());

	memset(m_tableHeaders[tableHandle]->m_name, 0, sizeof(char) * 256);
	strcpy(m_tableHeaders[tableHandle]->m_name, tableName);
	m_tableHeaders[tableHandle]->m_keyLength = keyLength;
	m_tableHeaders[tableHandle]->m_tablePosition = 0;
	m_tableHeaders[tableHandle]->m_tableSize = 0;
	m_tableHeaders[tableHandle]->m_allocationSize = 0;

	++m_dataStoreHeader.m_tableListSize;

	// Write table list
	SaveTableHeaders(m_dataStoreHeader, m_tableHeaders);

	// Reserve space for data store header
	SaveDataStoreHeader(&m_dataStoreHeader);

	// Return table handle
	return tableHandle;
}

BOOL CRuDataStore_Disk::DeleteTable(const char *tableName)
{
	return FALSE;
}

INT32 CRuDataStore_Disk::GetTableHandle(const char *tableName)
{
	for(INT32 i = 0; i < m_tableHeaders.Count(); ++i)
	{
		if(strcmpi(tableName, m_tableHeaders[i]->m_name) == 0)
		{
			return i;
		}
	}

	return -1;
}

INT32 CRuDataStore_Disk::GetTableKeyLength(INT32 tableHandle)
{
	// Ensure validity of handle
	assert(tableHandle >= 0 && tableHandle < m_tableHeaders.Count());

	return m_tableHeaders[tableHandle]->m_keyLength;
}

INT32 CRuDataStore_Disk::GetNumRecords(INT32 tableHandle)
{
	// Ensure validity of handle
	assert(tableHandle >= 0 && tableHandle < m_tableHeaders.Count());

	return m_tableHeaders[tableHandle]->m_streamHeaders.Count();
}

IRuStream *CRuDataStore_Disk::Retrieve(INT32 tableHandle, const BYTE *key)
{
	// Ensure validity of handle
	assert(tableHandle >= 0 && tableHandle < m_tableHeaders.Count());

	// Build a key structure
	KeyStruct keyStruct;
	keyStruct.m_key = key;
	keyStruct.m_keyLength = m_tableHeaders[tableHandle]->m_keyLength;

	if(m_tableHeaders[tableHandle]->m_streamHeaders.ContainsKey(keyStruct))
	{
		// Locate the stream header record
		StreamHeader &header = m_tableHeaders[tableHandle]->m_streamHeaders.Get(keyStruct);

		// Seek to position of the stream
		m_stream->Seek(header.m_streamPosition, ruSSM_Begin);

		// Read stream into memory
		CRuMemoryStream *memStream;
		memStream = ruNEW CRuMemoryStream();
		memStream->SetBufferSize(header.m_streamSize);
		memStream->SetStreamSize(header.m_streamSize);
		m_stream->Read(memStream->GetBuffer(), header.m_streamSize);

		return memStream;
	}

	return NULL;
}

IRuStream *CRuDataStore_Disk::RetrieveByIndex(INT32 tableHandle, INT32 index)
{
	// Ensure validity of handle
	assert(tableHandle >= 0 && tableHandle < m_tableHeaders.Count());

	if(index >= 0 && index < m_tableHeaders[tableHandle]->m_streamHeaders.Count())
	{
		// Locate the stream header record
		StreamHeader &header = m_tableHeaders[tableHandle]->m_streamHeaders.GetAtIndex(index);

		// Seek to position of the stream
		m_stream->Seek(header.m_streamPosition, ruSSM_Begin);

		if(header.m_streamSize > 1024 * 1024)
		{
			printf("\nWARNING: Ignored abnormally large stream chunk at %d, size %d\n", index, header.m_streamSize);

			return NULL;
		}

		// Read stream into memory
		CRuMemoryStream *memStream;
		memStream = ruNEW CRuMemoryStream();
		memStream->SetBufferSize(header.m_streamSize);
		memStream->SetStreamSize(header.m_streamSize);
		m_stream->Read(memStream->GetBuffer(), header.m_streamSize);

		return memStream;
	}

	return NULL;
}

BOOL CRuDataStore_Disk::RetrieveKeyByIndex(INT32 tableHandle, INT32 index, BYTE *keyOut)
{
	// Ensure validity of handle
	assert(tableHandle >= 0 && tableHandle < m_tableHeaders.Count());

	if(index >= 0 && index < m_tableHeaders[tableHandle]->m_streamHeaders.Count())
	{
		// Locate the stream header record
		StreamHeader &header = m_tableHeaders[tableHandle]->m_streamHeaders.GetAtIndex(index);

		// Copy key to output
		memcpy(keyOut, header.m_key, m_tableHeaders[tableHandle]->m_keyLength);

		return TRUE;
	}

	return FALSE;
}

BOOL CRuDataStore_Disk::Store(INT32 tableHandle, const BYTE *key, IRuStream *stream)
{
	StreamHeader streamHeader;

	// Ensure validity of handle
	assert(tableHandle >= 0 && tableHandle < m_tableHeaders.Count());

	if(m_readOnly)
	{
		return FALSE;
	}

	// Build a key structure
	KeyStruct keyStruct;
	keyStruct.m_key = key;
	keyStruct.m_keyLength = m_tableHeaders[tableHandle]->m_keyLength;

	if(m_tableHeaders[tableHandle]->m_streamHeaders.ContainsKey(keyStruct))
	{
		// Locate the stream header record
		streamHeader = m_tableHeaders[tableHandle]->m_streamHeaders.Get(keyStruct);

		if(streamHeader.m_allocationSize < (INT32) stream->GetStreamSize())
		{
			m_stream->Seek(0, ruSSM_End);
			streamHeader.m_streamPosition = m_stream->GetPosition();
			streamHeader.m_allocationSize = 0;
		}

		// Update key struct with the correct key pointer
		keyStruct.m_key = streamHeader.m_key;
	}
	else
	{
		// Set up stream header
		streamHeader.m_key = ruNEW BYTE [m_tableHeaders[tableHandle]->m_keyLength];
		memcpy(streamHeader.m_key, key, m_tableHeaders[tableHandle]->m_keyLength);

		m_stream->Seek(0, ruSSM_End);
		streamHeader.m_streamPosition = m_stream->GetPosition();
		streamHeader.m_allocationSize = 0;

		// Update key struct with the correct key pointer
		keyStruct.m_key = streamHeader.m_key;

		// Append stream header to table
		++m_tableHeaders[tableHandle]->m_tableSize;
		m_tableHeaders[tableHandle]->m_streamHeaders.Add(keyStruct, streamHeader);
	}

	// Update stream size
	streamHeader.m_streamSize = stream->GetStreamSize();

	// Output stream data
	stream->Seek(0, ruSSM_Begin);
	m_stream->Seek(streamHeader.m_streamPosition, ruSSM_Begin);
	do
	{
		memset(m_ioBuffer, 0, m_ioBufferSize);
		INT32 bytesRead = stream->Read(m_ioBuffer, m_ioBufferSize);

		if(bytesRead > 0)
		{
			// Update allocation size if necessary
			if(streamHeader.m_allocationSize < streamHeader.m_streamSize)
			{
				streamHeader.m_allocationSize += m_ioBufferSize;
			}

			// Always write in multiples of the io buffer size
			m_stream->Write(m_ioBuffer, m_ioBufferSize);
		}
		else
		{
			break;
		}
	} while(TRUE);

	// Update stream header
	m_tableHeaders[tableHandle]->m_streamHeaders.Set(keyStruct, streamHeader);

	// Update stream list
//	SaveStreamHeaders(m_tableHeaders, tableHandle);

	// Update table list
//	SaveTableHeaders(m_dataStoreHeader, m_tableHeaders);

	// Update data store header
//	SaveDataStoreHeader(&m_dataStoreHeader);

	return TRUE;
}

BOOL CRuDataStore_Disk::Delete(INT32 tableHandle, const BYTE *key)
{
	// Ensure validity of handle
	assert(tableHandle >= 0 && tableHandle < m_tableHeaders.Count());

	if(m_readOnly)
	{
		return FALSE;
	}

	// Build a key structure
	KeyStruct keyStruct;
	keyStruct.m_key = key;
	keyStruct.m_keyLength = m_tableHeaders[tableHandle]->m_keyLength;

	if(m_tableHeaders[tableHandle]->m_streamHeaders.ContainsKey(keyStruct))
	{
		// Remove the stream header record
		--m_tableHeaders[tableHandle]->m_tableSize;
		m_tableHeaders[tableHandle]->m_streamHeaders.Remove(keyStruct);

		// Update stream list
//		SaveStreamHeaders(m_tableHeaders, tableHandle);

		// Update table list
//		SaveTableHeaders(m_dataStoreHeader, m_tableHeaders);

		// Update data store header
//		SaveDataStoreHeader(&m_dataStoreHeader);

		return TRUE;
	}

	return FALSE;
}

BOOL CRuDataStore_Disk::DeleteAllRecords(INT32 tableHandle)
{
	// Ensure validity of handle
	assert(tableHandle >= 0 && tableHandle < m_tableHeaders.Count());

	if(m_readOnly)
	{
		return FALSE;
	}

	// Clear table
	m_tableHeaders[tableHandle]->m_tableSize = 0;
	m_tableHeaders[tableHandle]->m_streamHeaders.Clear();

	return TRUE;
}

BOOL CRuDataStore_Disk::Flush()
{
	if(m_readOnly)
	{
		return FALSE;
	}

	for(INT32 i = 0; i < m_tableHeaders.Count(); ++i)
	{
		// Update stream list
		SaveStreamHeaders(m_tableHeaders, i);
	}

	// Update table list
	SaveTableHeaders(m_dataStoreHeader, m_tableHeaders);

	// Update data store header
	SaveDataStoreHeader(&m_dataStoreHeader);

	return TRUE;
}

BOOL CRuDataStore_Disk::LoadDataStoreHeader(DataStoreHeader *headerOut)
{
	// Seek to beginning of data store stream
	m_stream->Seek(0, ruSSM_Begin);

	// Read & confirm GUID correctness
	CRuGUID classGUID;
	if(RuReadGUID(m_stream, classGUID) == FALSE)
	{
		return FALSE;
	}

	// Read data store header position
	m_stream->Read(&headerOut->m_tableListPosition, sizeof(INT32));
	m_stream->Read(&headerOut->m_tableListSize, sizeof(INT32));
	m_stream->Read(&headerOut->m_tableListAllocationSize, sizeof(INT32));

	return TRUE;
}

BOOL CRuDataStore_Disk::SaveDataStoreHeader(DataStoreHeader *header)
{
	// Seek to beginning of data store stream
	m_stream->Seek(0, ruSSM_Begin);

	// Write & confirm GUID correctness
	m_stream->Write(GetClassGUID().Data(), GetClassGUID().Size());

	// Write data store header position
	m_stream->Write(&header->m_tableListPosition, sizeof(INT32));
	m_stream->Write(&header->m_tableListSize, sizeof(INT32));
	m_stream->Write(&header->m_tableListAllocationSize, sizeof(INT32));

	return TRUE;
}

BOOL CRuDataStore_Disk::LoadTableHeaders(DataStoreHeader &dsHeader, CRuArrayList<TableHeader *> &tableHeaders)
{
	// Seek to beginning of table headers
	m_stream->Seek(dsHeader.m_tableListPosition, ruSSM_Begin);

	// Clear table headers array
	tableHeaders.Clear();

	for(INT32 i = 0; i < dsHeader.m_tableListSize; ++i)
	{
		// Append blank header to table list
		tableHeaders.Add(ruNEW TableHeader());

		// Read table properties
		m_stream->Read(tableHeaders[i]->m_name, sizeof(char) * 256);
		m_stream->Read(&tableHeaders[i]->m_keyLength, sizeof(INT32));
		m_stream->Read(&tableHeaders[i]->m_tablePosition, sizeof(INT32));
		m_stream->Read(&tableHeaders[i]->m_tableSize, sizeof(INT32));
		m_stream->Read(&tableHeaders[i]->m_allocationSize, sizeof(INT32));
	}

	return TRUE;
}

BOOL CRuDataStore_Disk::SaveTableHeaders(DataStoreHeader &dsHeader, CRuArrayList<TableHeader *> &tableHeaders)
{
	// If the current allocation size is too small, relocate the table list
	if(dsHeader.m_tableListAllocationSize < dsHeader.m_tableListSize || dsHeader.m_tableListAllocationSize == 0)
	{
		// Set table list position to the end of the stream
		dsHeader.m_tableListPosition = m_stream->GetStreamSize();

		// Round allocation size up to the nearest multiple of m_tableListPageSize
		while(dsHeader.m_tableListAllocationSize < dsHeader.m_tableListSize || dsHeader.m_tableListAllocationSize == 0)
		{
			dsHeader.m_tableListAllocationSize += m_tableListPageSize;
		}
	}

	// Seek to beginning of table headers
	m_stream->Seek(dsHeader.m_tableListPosition, ruSSM_Begin);

	// Write blank entries for the entire table list allocation section
	TableHeader dummyHeader;
	memset(&dummyHeader, 0, sizeof(TableHeader));
	for(INT32 i = 0; i < dsHeader.m_tableListAllocationSize; ++i)
	{
		m_stream->Write(dummyHeader.m_name, sizeof(char) * 256);
		m_stream->Write(&dummyHeader.m_keyLength, sizeof(INT32));
		m_stream->Write(&dummyHeader.m_tablePosition, sizeof(INT32));
		m_stream->Write(&dummyHeader.m_tableSize, sizeof(INT32));
		m_stream->Write(&dummyHeader.m_allocationSize, sizeof(INT32));
	}

	// Seek to beginning of table headers
	m_stream->Seek(dsHeader.m_tableListPosition, ruSSM_Begin);

	// Write occupied table headers
	for(INT32 i = 0; i < tableHeaders.Count(); ++i)
	{
		// Write table properties
		m_stream->Write(tableHeaders[i]->m_name, sizeof(char) * 256);
		m_stream->Write(&tableHeaders[i]->m_keyLength, sizeof(INT32));
		m_stream->Write(&tableHeaders[i]->m_tablePosition, sizeof(INT32));
		m_stream->Write(&tableHeaders[i]->m_tableSize, sizeof(INT32));
		m_stream->Write(&tableHeaders[i]->m_allocationSize, sizeof(INT32));
	}

	return TRUE;
}

BOOL CRuDataStore_Disk::LoadStreamHeaders(CRuArrayList<TableHeader *> &tableHeaders, INT32 tableHandle)
{
	// Seek to beginning of stream headers
	m_stream->Seek(tableHeaders[tableHandle]->m_tablePosition, ruSSM_Begin);

	for(INT32 i = 0; i < tableHeaders[tableHandle]->m_tableSize; ++i)
	{
		// Read stream header
		StreamHeader header;
		header.m_key = ruNEW BYTE [tableHeaders[tableHandle]->m_keyLength];
		m_stream->Read(header.m_key, tableHeaders[tableHandle]->m_keyLength);
		m_stream->Read(&header.m_streamPosition, sizeof(INT32));
		m_stream->Read(&header.m_streamSize, sizeof(INT32));
		m_stream->Read(&header.m_allocationSize, sizeof(INT32));

		// Build a key structure
		KeyStruct keyStruct;
		keyStruct.m_key = header.m_key;
		keyStruct.m_keyLength = tableHeaders[tableHandle]->m_keyLength;

		tableHeaders[tableHandle]->m_streamHeaders.Add(keyStruct, header);
	}

	// Balance the stream header tree
	tableHeaders[tableHandle]->m_streamHeaders.BalanceTree();

	return TRUE;
}

BOOL CRuDataStore_Disk::SaveStreamHeaders(CRuArrayList<TableHeader *> &tableHeaders, INT32 tableHandle)
{
	// If the current allocation size is too small, relocate the table
	if(tableHeaders[tableHandle]->m_allocationSize < tableHeaders[tableHandle]->m_tableSize)
	{
		// Set table list position to the end of the stream
		tableHeaders[tableHandle]->m_tablePosition = m_stream->GetStreamSize();

		// Round allocation size up to the nearest multiple of m_tablePageSize
		while(tableHeaders[tableHandle]->m_allocationSize < tableHeaders[tableHandle]->m_tableSize)
		{
			tableHeaders[tableHandle]->m_allocationSize += m_tablePageSize;
		}
	}

	// Reserve space for the entire table allocation section
	INT32 streamHeaderSize = sizeof(INT32) * 3 + tableHeaders[tableHandle]->m_keyLength;
	INT32 allocatedBlockSize = streamHeaderSize * tableHeaders[tableHandle]->m_allocationSize;

	// Store temporaries for the save stream headers callback
	m_callbackKeyLength = tableHeaders[tableHandle]->m_keyLength;

	m_tempStream.Seek(0, ruSSM_Begin);
	m_tempStream.SetStreamSize(allocatedBlockSize);

	// Balance the stream header tree
	tableHeaders[tableHandle]->m_streamHeaders.BalanceTree();

	// Iterate over all stream headers
	tableHeaders[tableHandle]->m_streamHeaders.ForAllElements(boost::bind(&CRuDataStore_Disk::SaveStreamHeadersCallback, this, _1, _2));

	// Seek to beginning of stream headers
	m_stream->Seek(tableHeaders[tableHandle]->m_tablePosition, ruSSM_Begin);

	// Write entire stream header block
	m_stream->Write(m_tempStream.GetBuffer(), m_tempStream.GetStreamSize());

	return TRUE;
}

BOOL CRuDataStore_Disk::SaveStreamHeadersCallback(const void *key, void *data)
{
	StreamHeader *header = (StreamHeader *) data;

	m_tempStream.Write(header->m_key, m_callbackKeyLength);
	m_tempStream.Write(&header->m_streamPosition, sizeof(INT32));
	m_tempStream.Write(&header->m_streamSize, sizeof(INT32));
	m_tempStream.Write(&header->m_allocationSize, sizeof(INT32));

	return TRUE;
}

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

#pragma managed
