#include "RenderEngine.h"

#include <gl/gl.h>

#include <algorithm>
#include <cstring>

RenderEngine::RenderEngine(
	RuntimeSnapshotProvider& runtimeSnapshotProvider,
	HealthTelemetry& healthTelemetry,
	HDC hdc,
	HGLRC hglrc,
	RenderEffectCallback renderEffect,
	ScreenshotCallback screenshotReady,
	PreviewPaintCallback previewPaint) :
	mRenderer(),
	mRenderPass(mRenderer),
	mRenderPipeline(mRenderer, runtimeSnapshotProvider, healthTelemetry, std::move(renderEffect), std::move(screenshotReady), std::move(previewPaint)),
	mShaderPrograms(mRenderer, runtimeSnapshotProvider),
	mHdc(hdc),
	mHglrc(hglrc),
	mFrameStateResolver(runtimeSnapshotProvider)
{
}

RenderEngine::~RenderEngine()
{
	StopRenderThread();
	if (!mResourcesDestroyed)
	{
		mRenderer.DestroyResources();
		mResourcesDestroyed = true;
	}
}

bool RenderEngine::StartRenderThread()
{
	if (mRenderThreadRunning)
		return true;

	{
		std::lock_guard<std::mutex> lock(mRenderThreadMutex);
		mRenderThreadStopping = false;
	}

	std::promise<bool> ready;
	std::future<bool> readyResult = ready.get_future();
	mRenderThread = std::thread(&RenderEngine::RenderThreadMain, this, std::move(ready));
	if (!readyResult.get())
	{
		if (mRenderThread.joinable())
			mRenderThread.join();
		return false;
	}

	return true;
}

void RenderEngine::StopRenderThread()
{
	if (mRenderThreadRunning)
	{
		InvokeOnRenderThread([this]() {
			if (!mResourcesDestroyed)
			{
				mRenderer.DestroyResources();
				mResourcesDestroyed = true;
			}
		});
	}

	{
		std::lock_guard<std::mutex> lock(mRenderThreadMutex);
		mRenderThreadStopping = true;
	}
	mRenderThreadCondition.notify_one();

	if (mRenderThread.joinable())
		mRenderThread.join();
}

void RenderEngine::RenderThreadMain(std::promise<bool> ready)
{
	mRenderThreadId = GetCurrentThreadId();
	if (!wglMakeCurrent(mHdc, mHglrc))
	{
		mRenderThreadId = 0;
		ready.set_value(false);
		return;
	}

	mRenderThreadRunning = true;
	ready.set_value(true);

	for (;;)
	{
		std::function<void()> task;
		{
			std::unique_lock<std::mutex> lock(mRenderThreadMutex);
			mRenderThreadCondition.wait(lock, [this]() {
				return mRenderThreadStopping || !mRenderThreadTasks.empty();
			});

			if (mRenderThreadStopping && mRenderThreadTasks.empty())
				break;

			task = std::move(mRenderThreadTasks.front());
			mRenderThreadTasks.pop();
		}

		try
		{
			task();
		}
		catch (...)
		{
			OutputDebugStringA("Render thread task failed with an unhandled exception.\n");
		}
	}

	wglMakeCurrent(NULL, NULL);
	mRenderThreadRunning = false;
	mRenderThreadId = 0;
}

void RenderEngine::ReportRenderThreadRequestFailure(const char* operationName, const char* reason)
{
	std::ostringstream message;
	message << "Render thread request failed";
	if (operationName && operationName[0] != '\0')
		message << " [" << operationName << "]";
	if (reason && reason[0] != '\0')
		message << ": " << reason;
	message << ".\n";
	OutputDebugStringA(message.str().c_str());
}

bool RenderEngine::IsRenderThreadAccessExpected() const
{
	return !mRenderThreadRunning || GetCurrentThreadId() == mRenderThreadId;
}

void RenderEngine::ReportWrongThreadRenderAccess(const char* operationName) const
{
	if (IsRenderThreadAccessExpected())
		return;

	std::ostringstream message;
	message << "Wrong-thread render access detected";
	if (operationName && operationName[0] != '\0')
		message << " [" << operationName << "]";
	message << ".\n";
	OutputDebugStringA(message.str().c_str());
}

bool RenderEngine::CompileDecodeShader(int errorMessageSize, char* errorMessage)
{
	return InvokeOnRenderThread([this, errorMessageSize, errorMessage]() {
		return mShaderPrograms.CompileDecodeShader(errorMessageSize, errorMessage);
	});
}

bool RenderEngine::CompileOutputPackShader(int errorMessageSize, char* errorMessage)
{
	return InvokeOnRenderThread([this, errorMessageSize, errorMessage]() {
		return mShaderPrograms.CompileOutputPackShader(errorMessageSize, errorMessage);
	});
}

bool RenderEngine::InitializeResources(
	unsigned inputFrameWidth,
	unsigned inputFrameHeight,
	unsigned captureTextureWidth,
	unsigned outputFrameWidth,
	unsigned outputFrameHeight,
	unsigned outputPackTextureWidth,
	std::string& error)
{
	return InvokeOnRenderThread([this, inputFrameWidth, inputFrameHeight, captureTextureWidth, outputFrameWidth, outputFrameHeight, outputPackTextureWidth, &error]() {
		return mRenderer.InitializeResources(
			inputFrameWidth,
			inputFrameHeight,
			captureTextureWidth,
			outputFrameWidth,
			outputFrameHeight,
			outputPackTextureWidth,
			error);
	});
}

bool RenderEngine::CompileLayerPrograms(unsigned inputFrameWidth, unsigned inputFrameHeight, int errorMessageSize, char* errorMessage)
{
	return InvokeOnRenderThread([this, inputFrameWidth, inputFrameHeight, errorMessageSize, errorMessage]() {
		return mShaderPrograms.CompileLayerPrograms(inputFrameWidth, inputFrameHeight, errorMessageSize, errorMessage);
	});
}

bool RenderEngine::CommitPreparedLayerPrograms(const PreparedShaderBuild& preparedBuild, unsigned inputFrameWidth, unsigned inputFrameHeight, int errorMessageSize, char* errorMessage)
{
	return InvokeOnRenderThread([this, &preparedBuild, inputFrameWidth, inputFrameHeight, errorMessageSize, errorMessage]() {
		return mShaderPrograms.CommitPreparedLayerPrograms(preparedBuild, inputFrameWidth, inputFrameHeight, errorMessageSize, errorMessage);
	});
}

bool RenderEngine::ApplyPreparedShaderBuild(
	const PreparedShaderBuild& preparedBuild,
	unsigned inputFrameWidth,
	unsigned inputFrameHeight,
	bool preserveFeedbackState,
	int errorMessageSize,
	char* errorMessage)
{
	if (!CommitPreparedLayerPrograms(preparedBuild, inputFrameWidth, inputFrameHeight, errorMessageSize, errorMessage))
		return false;

	mFrameStateResolver.StoreCommittedSnapshot(preparedBuild.renderSnapshot, mShaderPrograms.CommittedLayerStates());
	ResetTemporalHistoryState();
	if (!preserveFeedbackState)
		ResetShaderFeedbackState();
	return true;
}

void RenderEngine::ResetTemporalHistoryState()
{
	InvokeOnRenderThread([this]() {
		mRenderCommandQueue.RequestRenderReset(RenderCommandResetScope::TemporalHistoryOnly);
		ProcessRenderResetCommandsOnRenderThread();
	});
}

void RenderEngine::ResetShaderFeedbackState()
{
	InvokeOnRenderThread([this]() {
		mRenderCommandQueue.RequestRenderReset(RenderCommandResetScope::ShaderFeedbackOnly);
		ProcessRenderResetCommandsOnRenderThread();
	});
}

void RenderEngine::ApplyRuntimeCoordinatorRenderReset(RuntimeCoordinatorRenderResetScope resetScope)
{
	InvokeOnRenderThread([this, resetScope]() {
		switch (resetScope)
		{
		case RuntimeCoordinatorRenderResetScope::TemporalHistoryOnly:
			mRenderCommandQueue.RequestRenderReset(RenderCommandResetScope::TemporalHistoryOnly);
			ProcessRenderResetCommandsOnRenderThread();
			break;
		case RuntimeCoordinatorRenderResetScope::TemporalHistoryAndFeedback:
			mRenderCommandQueue.RequestRenderReset(RenderCommandResetScope::TemporalHistoryAndFeedback);
			ProcessRenderResetCommandsOnRenderThread();
			break;
		case RuntimeCoordinatorRenderResetScope::None:
		default:
			break;
		}
	});
}

void RenderEngine::ResetTemporalHistoryStateOnRenderThread()
{
	ReportWrongThreadRenderAccess("reset-temporal-history");
	mShaderPrograms.ResetTemporalHistoryState();
}

void RenderEngine::ResetShaderFeedbackStateOnRenderThread()
{
	ReportWrongThreadRenderAccess("reset-shader-feedback");
	mShaderPrograms.ResetShaderFeedbackState();
}

void RenderEngine::ApplyRenderResetOnRenderThread(RenderCommandResetScope resetScope)
{
	switch (resetScope)
	{
	case RenderCommandResetScope::ShaderFeedbackOnly:
		ResetShaderFeedbackStateOnRenderThread();
		break;
	case RenderCommandResetScope::TemporalHistoryOnly:
		ResetTemporalHistoryStateOnRenderThread();
		break;
	case RenderCommandResetScope::TemporalHistoryAndFeedback:
		ResetTemporalHistoryStateOnRenderThread();
		ResetShaderFeedbackStateOnRenderThread();
		break;
	case RenderCommandResetScope::None:
	default:
		break;
	}
}

void RenderEngine::ProcessRenderResetCommandsOnRenderThread()
{
	RenderCommandResetScope resetScope = RenderCommandResetScope::None;
	while (mRenderCommandQueue.TryTakeRenderReset(resetScope))
		ApplyRenderResetOnRenderThread(resetScope);
}

void RenderEngine::EnqueuePreviewPresentWake()
{
	if (!mRenderThreadRunning)
		return;

	bool shouldNotify = false;
	{
		std::lock_guard<std::mutex> lock(mRenderThreadMutex);
		if (!mRenderThreadStopping && !mPreviewPresentWakePending)
		{
			mPreviewPresentWakePending = true;
			mRenderThreadTasks.push([this]() {
				{
					std::lock_guard<std::mutex> lock(mRenderThreadMutex);
					mPreviewPresentWakePending = false;
				}
				ProcessPreviewPresentCommandsOnRenderThread();
			});
			shouldNotify = true;
		}
	}

	if (shouldNotify)
		mRenderThreadCondition.notify_one();
}

void RenderEngine::ProcessPreviewPresentCommandsOnRenderThread()
{
	RenderPreviewPresentRequest request;
	if (mRenderCommandQueue.TryTakePreviewPresent(request))
		PresentPreviewOnRenderThread(request.outputFrameWidth, request.outputFrameHeight);
}

void RenderEngine::EnqueueInputUploadWake()
{
	if (!mRenderThreadRunning || GetCurrentThreadId() == mRenderThreadId)
		return;

	bool shouldNotify = false;
	{
		std::lock_guard<std::mutex> lock(mRenderThreadMutex);
		if (!mRenderThreadStopping && !mInputUploadWakePending)
		{
			mInputUploadWakePending = true;
			mRenderThreadTasks.push([this]() {
				{
					std::lock_guard<std::mutex> lock(mRenderThreadMutex);
					mInputUploadWakePending = false;
				}
				ProcessInputUploadCommandsOnRenderThread();
			});
			shouldNotify = true;
		}
	}

	if (shouldNotify)
		mRenderThreadCondition.notify_one();
}

void RenderEngine::ProcessInputUploadCommandsOnRenderThread()
{
	RenderInputUploadRequest request;
	while (mRenderCommandQueue.TryTakeInputUpload(request))
	{
		if (request.ownedBytes.empty())
			continue;

		request.inputFrame.bytes = request.ownedBytes.data();
		UploadInputFrameOnRenderThread(request.inputFrame, request.videoState);
	}
}

void RenderEngine::EnqueueScreenshotCaptureWake()
{
	if (!mRenderThreadRunning || GetCurrentThreadId() == mRenderThreadId)
		return;

	bool shouldNotify = false;
	{
		std::lock_guard<std::mutex> lock(mRenderThreadMutex);
		if (!mRenderThreadStopping && !mScreenshotCaptureWakePending)
		{
			mScreenshotCaptureWakePending = true;
			mRenderThreadTasks.push([this]() {
				{
					std::lock_guard<std::mutex> lock(mRenderThreadMutex);
					mScreenshotCaptureWakePending = false;
				}
				ProcessScreenshotCaptureCommandsOnRenderThread();
			});
			shouldNotify = true;
		}
	}

	if (shouldNotify)
		mRenderThreadCondition.notify_one();
}

void RenderEngine::ProcessScreenshotCaptureCommandsOnRenderThread()
{
	RenderScreenshotCaptureRequest request;
	ScreenshotCaptureCallback completion;
	{
		std::lock_guard<std::mutex> lock(mRenderThreadMutex);
		completion = mScreenshotCaptureCompletion;
	}

	while (mRenderCommandQueue.TryTakeScreenshotCapture(request))
	{
		if (!completion)
			continue;

		std::vector<unsigned char> topDownPixels;
		if (CaptureOutputFrameRgbaTopDownOnRenderThread(request.width, request.height, topDownPixels))
			completion(request.width, request.height, std::move(topDownPixels));
	}
}

void RenderEngine::ClearOscOverlayState()
{
	InvokeOnRenderThread([this]() {
		mRuntimeLiveState.Clear();
	});
}

void RenderEngine::ClearOscOverlayStateForLayerKey(const std::string& layerKey)
{
	InvokeOnRenderThread([this, layerKey]() {
		mRuntimeLiveState.ClearForLayerKey(layerKey);
	});
}

void RenderEngine::UpdateOscOverlayState(
	const std::vector<OscOverlayUpdate>& updates,
	const std::vector<OscOverlayCommitCompletion>& completedCommits)
{
	std::vector<RuntimeLiveOscCommitCompletion> liveCompletions;
	liveCompletions.reserve(completedCommits.size());
	for (const OscOverlayCommitCompletion& completedCommit : completedCommits)
		liveCompletions.push_back({ completedCommit.routeKey, completedCommit.generation });
	mRuntimeLiveState.ApplyOscCommitCompletions(liveCompletions);

	std::vector<RuntimeLiveOscUpdate> liveUpdates;
	liveUpdates.reserve(updates.size());
	for (const OscOverlayUpdate& update : updates)
		liveUpdates.push_back({ update.routeKey, update.layerKey, update.parameterKey, update.targetValue });
	mRuntimeLiveState.ApplyOscUpdates(liveUpdates);
}

void RenderEngine::ResizeView(int width, int height)
{
	InvokeOnRenderThread([this, width, height]() {
		mRenderer.ResizeView(width, height);
	});
}

bool RenderEngine::TryPresentPreview(bool force, unsigned previewFps, unsigned outputFrameWidth, unsigned outputFrameHeight)
{
	if (!force)
	{
		if (previewFps == 0)
			return false;

		const auto now = std::chrono::steady_clock::now();
		const auto minimumInterval = std::chrono::microseconds(1000000 / (previewFps == 0 ? 1u : previewFps));
		if (mLastPreviewPresentTime != std::chrono::steady_clock::time_point() &&
			now - mLastPreviewPresentTime < minimumInterval)
		{
			return false;
		}
	}

	if (mRenderThreadRunning)
	{
		mRenderCommandQueue.RequestPreviewPresent({ outputFrameWidth, outputFrameHeight });
		EnqueuePreviewPresentWake();
		return true;
	}

	ReportRenderThreadRequestFailure("preview-present", "render thread is not running");
	return false;
}

bool RenderEngine::PresentPreviewOnRenderThread(unsigned outputFrameWidth, unsigned outputFrameHeight)
{
	ReportWrongThreadRenderAccess("preview-present");
	mRenderer.PresentToWindow(mHdc, outputFrameWidth, outputFrameHeight);
	mLastPreviewPresentTime = std::chrono::steady_clock::now();
	return true;
}

bool RenderEngine::RequestScreenshotCapture(unsigned width, unsigned height, ScreenshotCaptureCallback completion)
{
	if (width == 0 || height == 0 || !completion)
		return false;
	if (!mRenderThreadRunning)
		return false;

	{
		std::lock_guard<std::mutex> lock(mRenderThreadMutex);
		mScreenshotCaptureCompletion = std::move(completion);
	}
	mRenderCommandQueue.RequestScreenshotCapture({ width, height });
	EnqueueScreenshotCaptureWake();
	return true;
}

bool RenderEngine::QueueInputFrame(const VideoIOFrame& inputFrame, const VideoIOState& videoState)
{
	if (inputFrame.hasNoInputSource || inputFrame.bytes == nullptr)
		return true;
	if (inputFrame.rowBytes <= 0 || inputFrame.height == 0)
		return false;

	const std::size_t byteCount = static_cast<std::size_t>(inputFrame.rowBytes) * inputFrame.height;
	RenderInputUploadRequest request;
	request.inputFrame = inputFrame;
	request.videoState = videoState;
	request.ownedBytes.resize(byteCount);
	std::memcpy(request.ownedBytes.data(), inputFrame.bytes, byteCount);
	request.inputFrame.bytes = nullptr;

	mRenderCommandQueue.RequestInputUpload(request);
	EnqueueInputUploadWake();
	return true;
}

bool RenderEngine::UploadInputFrameOnRenderThread(const VideoIOFrame& inputFrame, const VideoIOState& videoState)
{
	ReportWrongThreadRenderAccess("input-upload");
	const long textureSize = inputFrame.rowBytes * static_cast<long>(inputFrame.height);
	glPixelStorei(GL_UNPACK_ALIGNMENT, 4);
	glPixelStorei(GL_UNPACK_ROW_LENGTH, 0);
	glBindBuffer(GL_PIXEL_UNPACK_BUFFER, mRenderer.TextureUploadBuffer());
	glBufferData(GL_PIXEL_UNPACK_BUFFER, textureSize, inputFrame.bytes, GL_DYNAMIC_DRAW);
	glBindTexture(GL_TEXTURE_2D, mRenderer.CaptureTexture());
	if (inputFrame.pixelFormat == VideoIOPixelFormat::V210)
		glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, videoState.captureTextureWidth, videoState.inputFrameSize.height, GL_RGBA, GL_UNSIGNED_BYTE, NULL);
	else
		glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, videoState.captureTextureWidth, videoState.inputFrameSize.height, GL_BGRA, GL_UNSIGNED_INT_8_8_8_8_REV, NULL);
	glBindTexture(GL_TEXTURE_2D, 0);
	glBindBuffer(GL_PIXEL_UNPACK_BUFFER, 0);

	return true;
}

bool RenderEngine::RequestOutputFrame(const RenderPipelineFrameContext& context, VideoIOOutputFrame& outputFrame)
{
	if (mRenderThreadRunning)
	{
		return TryInvokeOnRenderThread("output-render", [this, &context, &outputFrame]() {
			mRenderCommandQueue.RequestOutputFrame({ context.videoState, context.completion });
			RenderOutputFrameRequest request;
			return mRenderCommandQueue.TryTakeOutputFrame(request) &&
				RenderOutputFrameOnRenderThread({ request.videoState, request.completion }, outputFrame);
		});
	}

	ReportRenderThreadRequestFailure("output-render", "render thread is not running");
	return false;
}

bool RenderEngine::RenderOutputFrameOnRenderThread(const RenderPipelineFrameContext& context, VideoIOOutputFrame& outputFrame)
{
	ReportWrongThreadRenderAccess("output-render");
	ProcessRenderResetCommandsOnRenderThread();
	ProcessInputUploadCommandsOnRenderThread();
	return mRenderPipeline.RenderFrame(context, outputFrame);
}

bool RenderEngine::ResolveRenderFrameState(
	const RenderFrameInput& input,
	std::vector<OscOverlayCommitRequest>* commitRequests,
	RenderFrameState& frameState)
{
	std::vector<RuntimeLiveOscCommitRequest> liveCommitRequests;
	const bool resolved = mFrameStateResolver.Resolve(
		input,
		mShaderPrograms.CommittedLayerStates(),
		mRuntimeLiveState,
		commitRequests ? &liveCommitRequests : nullptr,
		frameState);

	if (commitRequests)
	{
		for (const RuntimeLiveOscCommitRequest& request : liveCommitRequests)
			commitRequests->push_back({ request.routeKey, request.layerKey, request.parameterKey, request.value, request.generation });
	}
	return resolved;
}

void RenderEngine::RenderPreparedFrame(const RenderFrameState& frameState)
{
	RenderLayerStack(
		frameState.hasInputSource,
		frameState.layerStates,
		frameState.inputFrameWidth,
		frameState.inputFrameHeight,
		frameState.captureTextureWidth,
		frameState.inputPixelFormat,
		frameState.historyCap);
}

void RenderEngine::RenderLayerStack(
	bool hasInputSource,
	const std::vector<RuntimeRenderState>& layerStates,
	unsigned inputFrameWidth,
	unsigned inputFrameHeight,
	unsigned captureTextureWidth,
	VideoIOPixelFormat inputPixelFormat,
	unsigned historyCap)
{
	ReportWrongThreadRenderAccess("render-layer-stack");
	mRenderPass.Render(
		hasInputSource,
		layerStates,
		inputFrameWidth,
		inputFrameHeight,
		captureTextureWidth,
		inputPixelFormat,
		historyCap,
		[this](const RuntimeRenderState& state, OpenGLRenderer::LayerProgram::TextBinding& textBinding, std::string& error) {
			return mShaderPrograms.UpdateTextBindingTexture(state, textBinding, error);
		},
		[this](const RuntimeRenderState& state, unsigned availableSourceHistoryLength, unsigned availableTemporalHistoryLength, bool feedbackAvailable) {
			return mShaderPrograms.UpdateGlobalParamsBuffer(state, availableSourceHistoryLength, availableTemporalHistoryLength, feedbackAvailable);
		});
}

bool RenderEngine::ReadOutputFrameRgbaOnRenderThread(unsigned width, unsigned height, std::vector<unsigned char>& bottomUpPixels)
{
	ReportWrongThreadRenderAccess("read-output-frame-rgba");
	if (width == 0 || height == 0)
		return false;

	bottomUpPixels.resize(static_cast<std::size_t>(width) * height * 4);
	glBindFramebuffer(GL_READ_FRAMEBUFFER, mRenderer.OutputFramebuffer());
	glReadBuffer(GL_COLOR_ATTACHMENT0);
	glPixelStorei(GL_PACK_ALIGNMENT, 1);
	glPixelStorei(GL_PACK_ROW_LENGTH, 0);
	glReadPixels(0, 0, width, height, GL_RGBA, GL_UNSIGNED_BYTE, bottomUpPixels.data());
	glPixelStorei(GL_PACK_ALIGNMENT, 4);

	return true;
}

bool RenderEngine::CaptureOutputFrameRgbaTopDownOnRenderThread(unsigned width, unsigned height, std::vector<unsigned char>& topDownPixels)
{
	std::vector<unsigned char> bottomUpPixels;
	if (!ReadOutputFrameRgbaOnRenderThread(width, height, bottomUpPixels))
		return false;

	topDownPixels.resize(bottomUpPixels.size());
	const std::size_t rowBytes = static_cast<std::size_t>(width) * 4;
	for (unsigned y = 0; y < height; ++y)
	{
		const unsigned sourceY = height - 1 - y;
		std::copy(
			bottomUpPixels.begin() + static_cast<std::ptrdiff_t>(sourceY * rowBytes),
			bottomUpPixels.begin() + static_cast<std::ptrdiff_t>((sourceY + 1) * rowBytes),
			topDownPixels.begin() + static_cast<std::ptrdiff_t>(y * rowBytes));
	}

	return true;
}