#include "HeroAsyncWriter.h"

/** [WRITER THREAD] Flushes the archive and reset the flush timer. */
void FHeroAsyncWriter::FlushArchiveAndResetTimer()
{
	// This should be the one and only place where we flush because we want the flush to happen only on the 
	// async writer thread (if threading is enabled)
	Ar.Flush();
	LastArchiveFlushTime = FPlatformTime::Seconds();
}

/** [WRITER THREAD] Serialize the contents of the ring buffer to disk */
void FHeroAsyncWriter::SerializeBufferToArchive()
{
	while (SerializeRequestCounter.GetValue() > 0)
	{
		// Grab a local copy of the end pos. It's ok if it changes on the client thread later on.
		// We won't be modifying it anyway and will later serialize new data in the next iteration.
		// Here we only serialize what we know exists at the beginning of this function.
		int32 ThisThreadStartPos = BufferStartPos.Load(EMemoryOrder::Relaxed);
		int32 ThisThreadEndPos   = BufferEndPos  .Load(EMemoryOrder::Relaxed);

		if (ThisThreadEndPos >= ThisThreadStartPos)
		{
			Ar.Serialize(Buffer.GetData() + ThisThreadStartPos, ThisThreadEndPos - ThisThreadStartPos);
		}
		else
		{
			// Data is wrapped around the ring buffer
			Ar.Serialize(Buffer.GetData() + ThisThreadStartPos, Buffer.Num() - ThisThreadStartPos);
			Ar.Serialize(Buffer.GetData(), ThisThreadEndPos);
		}

		// Modify the start pos. Only the worker thread modifies this value so it's ok to not guard it with a critical section.
		BufferStartPos = ThisThreadEndPos;

		// Decrement the request counter, we now know we serialized at least one request.
		// We might have serialized more requests but it's irrelevant, the counter will go down to 0 eventually
		SerializeRequestCounter.Decrement();

		// Flush the archive periodically if running on a separate thread
		if (Thread)
		{
			if ((FPlatformTime::Seconds() - LastArchiveFlushTime) > 0.2f )
			{
				FlushArchiveAndResetTimer();
			}
		}
		// If no threading is available or when we explicitly requested flush (see FlushBuffer), flush immediately after writing.
		// In some rare cases we may flush twice (see above) but that's ok. We need a clear division between flushing because of the timer
		// and force flush on demand.
		if (WantsArchiveFlush.GetValue() > 0)
		{
			FlushArchiveAndResetTimer();
			int32 FlushCount = WantsArchiveFlush.Decrement();
			check(FlushCount >= 0);
		}
	}
}

/** [CLIENT THREAD] Flush the memory buffer (doesn't force the archive to flush). Can only be used from inside of BufferPosCritical lock. */
void FHeroAsyncWriter::FlushBuffer()
{
	SerializeRequestCounter.Increment();
	if (!Thread)
	{
		SerializeBufferToArchive();
	}
	while (SerializeRequestCounter.GetValue() != 0)
	{
		FPlatformProcess::SleepNoStats(0);
	}
	// Make sure there's been no unexpected concurrency
	check(SerializeRequestCounter.GetValue() == 0);
}

FHeroAsyncWriter::FHeroAsyncWriter(FArchive& InAr)
	: Thread(nullptr)
	, Ar(InAr)
	, BufferStartPos(0)
	, BufferEndPos(0)
	, LastArchiveFlushTime(0.0)
{
	Buffer.AddUninitialized(InitialBufferSize);
	

	if (FPlatformProcess::SupportsMultithreading())
	{
		FString WriterName = FString::Printf(TEXT("FHeroAsyncWriter_%s"), *FPaths::GetBaseFilename(Ar.GetArchiveName()));
		FPlatformAtomics::InterlockedExchangePtr((void**)&Thread, FRunnableThread::Create(this, *WriterName, 0, TPri_BelowNormal));
	}
}

FHeroAsyncWriter::~FHeroAsyncWriter()
{
	FScopeLock WriteLock(&BufferPosCritical);
	WantsArchiveFlush.Increment();
	SerializeRequestCounter.Increment();
	
	SerializeBufferToArchive();
	
	//Flush();
	delete Thread;
	Thread = nullptr;
}

/** [CLIENT THREAD] Serialize data to buffer that will later be saved to disk by the async thread */
void FHeroAsyncWriter::Serialize(void* InData, int64 Length)
{
	if (!InData || Length <= 0)
	{
		return;
	}

	const uint8* Data = (uint8*)InData;

	FScopeLock WriteLock(&BufferPosCritical);

	const int32 ThisThreadEndPos = BufferEndPos.Load(EMemoryOrder::Relaxed);

	// Store the local copy of the current buffer start pos. It may get moved by the worker thread but we don't
	// care about it too much because we only modify BufferEndPos. Copy should be atomic enough. We only use it
	// for checking the remaining space in the buffer so underestimating is ok.
	{
		const int32 ThisThreadStartPos = BufferStartPos.Load(EMemoryOrder::Relaxed);
		// Calculate the remaining size in the ring buffer
		const int32 BufferFreeSize = ThisThreadStartPos <= ThisThreadEndPos ? (Buffer.Num() - ThisThreadEndPos + ThisThreadStartPos) : (ThisThreadStartPos - ThisThreadEndPos);
		// Make sure the buffer is BIGGER than we require otherwise we may calculate the wrong (0) buffer EndPos for StartPos = 0 and Length = Buffer.Num()
		if (BufferFreeSize <= Length)
		{
			// Force the async thread to call SerializeBufferToArchive even if it's currently empty
			FlushBuffer();

			// Resize the buffer if needed
			if (Length >= Buffer.Num())
			{
				// Keep the buffer bigger than we require so that % Buffer.Num() does not return 0 for Lengths = Buffer.Num()
				Buffer.SetNumUninitialized((int32)(Length + 1));
			}
		}
	}

	// We now know there's enough space in the buffer to copy data
	const int32 WritePos = ThisThreadEndPos;
	if ((WritePos + Length) <= Buffer.Num())
	{
		// Copy straight into the ring buffer
		FMemory::Memcpy(Buffer.GetData() + WritePos, Data, Length);
	}
	else
	{
		// Wrap around the ring buffer
		int32 BufferSizeToEnd = Buffer.Num() - WritePos;
		FMemory::Memcpy(Buffer.GetData() + WritePos, Data, BufferSizeToEnd);
		FMemory::Memcpy(Buffer.GetData(), Data + BufferSizeToEnd, Length - BufferSizeToEnd);
	}

	// Update the end position and let the async thread know we need to write to disk
	BufferEndPos = (ThisThreadEndPos + Length) % Buffer.Num();
	SerializeRequestCounter.Increment();

	// No async thread? Serialize now.
	if (!Thread)
	{
		SerializeBufferToArchive();
	}
}

/** Flush all buffers to disk */
void FHeroAsyncWriter::Flush()
{
	FScopeLock WriteLock(&BufferPosCritical);
	WantsArchiveFlush.Increment();
	FlushBuffer();
}

//~ Begin FRunnable Interface.
bool FHeroAsyncWriter::Init()
{
	return true;
}
	
uint32 FHeroAsyncWriter::Run()
{
	while (StopTaskCounter.GetValue() == 0)
	{
		if (SerializeRequestCounter.GetValue() > 0)
		{
			SerializeBufferToArchive();
		}
		else if ((FPlatformTime::Seconds() - LastArchiveFlushTime) > 0.2f )
		{
			FlushArchiveAndResetTimer();
		}
		else
		{
			FPlatformProcess::SleepNoStats(0.01f);
		}
	}
	return 0;
}

void FHeroAsyncWriter::Stop()
{
	StopTaskCounter.Increment();
}