#include "RuntimeShaderPrepareWorker.h"

#include "../../platform/HiddenGlWindow.h"
#include "RuntimeShaderRenderer.h"

#include <algorithm>
#include <chrono>
#include <functional>
#include <utility>

RuntimeShaderPrepareWorker::~RuntimeShaderPrepareWorker()
{
	Stop();
}

bool RuntimeShaderPrepareWorker::Start(std::unique_ptr<HiddenGlWindow> sharedWindow, std::string& error)
{
	if (mThread.joinable())
		return true;
	if (!sharedWindow || sharedWindow->DeviceContext() == nullptr || sharedWindow->Context() == nullptr)
	{
		error = "Runtime shader prepare worker needs an existing shared GL context.";
		return false;
	}

	mWindow = std::move(sharedWindow);
	mStopping.store(false, std::memory_order_release);
	mStarted.store(false, std::memory_order_release);
	{
		std::lock_guard<std::mutex> lock(mMutex);
		mStartupReady = false;
		mStartupError.clear();
	}
	mThread = std::thread([this]() { ThreadMain(); });

	std::unique_lock<std::mutex> lock(mMutex);
	if (!mStartupCondition.wait_for(lock, std::chrono::seconds(3), [this]() {
		return mStartupReady || !mStartupError.empty();
	}))
	{
		error = "Timed out starting runtime shader prepare worker.";
		lock.unlock();
		Stop();
		return false;
	}
	if (!mStartupError.empty())
	{
		error = mStartupError;
		lock.unlock();
		Stop();
		return false;
	}
	return true;
}

void RuntimeShaderPrepareWorker::Stop()
{
	mStopping.store(true, std::memory_order_release);
	mCondition.notify_all();
	if (mThread.joinable())
		mThread.join();

	std::deque<RuntimePreparedShaderProgram> completed;
	{
		std::lock_guard<std::mutex> lock(mMutex);
		mRequests.clear();
		completed.swap(mCompleted);
	}
	for (RuntimePreparedShaderProgram& program : completed)
		program.ReleaseGl();
	mWindow.reset();
	mStarted.store(false, std::memory_order_release);
}

void RuntimeShaderPrepareWorker::Submit(const std::vector<RenderCadenceCompositor::RuntimeRenderLayerModel>& layers)
{
	std::lock_guard<std::mutex> lock(mMutex);
	for (const RenderCadenceCompositor::RuntimeRenderLayerModel& layer : layers)
	{
		if (layer.artifact.passes.empty() && layer.artifact.fragmentShaderSource.empty())
			continue;

		std::vector<RuntimeShaderPassArtifact> passes = layer.artifact.passes;
		if (passes.empty())
		{
			RuntimeShaderPassArtifact pass;
			pass.passId = "main";
			pass.fragmentShaderSource = layer.artifact.fragmentShaderSource;
			pass.outputName = "layerOutput";
			passes.push_back(std::move(pass));
		}

		auto sameLayer = [&layer](const PrepareRequest& existing) {
			return existing.layerId == layer.id;
		};
		mRequests.erase(std::remove_if(mRequests.begin(), mRequests.end(), sameLayer), mRequests.end());

		for (const RuntimeShaderPassArtifact& pass : passes)
		{
			PrepareRequest request;
			request.layerId = layer.id;
			request.shaderId = layer.shaderId;
			request.passId = pass.passId;
			request.sourceFingerprint = Fingerprint(layer.artifact);
			request.artifact = layer.artifact;
			request.passArtifact = pass;
			mRequests.push_back(std::move(request));
		}
	}
	mCondition.notify_one();
}

bool RuntimeShaderPrepareWorker::TryConsume(RuntimePreparedShaderProgram& preparedProgram)
{
	std::lock_guard<std::mutex> lock(mMutex);
	if (mCompleted.empty())
		return false;

	preparedProgram = std::move(mCompleted.front());
	mCompleted.pop_front();
	return true;
}

void RuntimeShaderPrepareWorker::ThreadMain()
{
	if (!mWindow || !mWindow->MakeCurrent())
	{
		std::lock_guard<std::mutex> lock(mMutex);
		mStartupError = "Runtime shader prepare worker could not make shared GL context current.";
		mStartupCondition.notify_all();
		return;
	}
	{
		std::lock_guard<std::mutex> lock(mMutex);
		mStartupReady = true;
	}
	mStarted.store(true, std::memory_order_release);
	mStartupCondition.notify_all();

	while (!mStopping.load(std::memory_order_acquire))
	{
		PrepareRequest request;
		{
			std::unique_lock<std::mutex> lock(mMutex);
			mCondition.wait(lock, [this]() {
				return mStopping.load(std::memory_order_acquire) || !mRequests.empty();
			});
			if (mStopping.load(std::memory_order_acquire))
				break;
			request = std::move(mRequests.front());
			mRequests.pop_front();
		}

		RuntimePreparedShaderProgram preparedProgram;
		RuntimeShaderRenderer::BuildPreparedPassProgram(
			request.layerId,
			request.sourceFingerprint,
			request.artifact,
			request.passArtifact,
			preparedProgram);
		glFlush();

		std::lock_guard<std::mutex> lock(mMutex);
		mCompleted.push_back(std::move(preparedProgram));
	}

	mWindow->ClearCurrent();
}

std::string RuntimeShaderPrepareWorker::Fingerprint(const RuntimeShaderArtifact& artifact)
{
	const std::hash<std::string> hasher;
	std::string source;
	for (const RuntimeShaderPassArtifact& pass : artifact.passes)
		source += pass.passId + ":" + pass.outputName + ":" + pass.fragmentShaderSource + "\n";
	if (source.empty())
		source = artifact.fragmentShaderSource;
	return artifact.shaderId + ":" + std::to_string(source.size()) + ":" + std::to_string(hasher(source));
}
