#include "RuntimeLayerController.h"

#include "../logging/Logger.h"

#include <algorithm>

namespace RenderCadenceCompositor
{
namespace
{
constexpr std::size_t kMaxConcurrentRuntimeShaderBuilds = 2;
}

void RuntimeLayerController::StartLayerShaderBuild(const std::string& layerId, const std::string& shaderId, bool preserveExistingRenderArtifact)
{
	CleanupRetiredShaderBuilds();
	RetireLayerShaderBuild(layerId);

	{
		std::lock_guard<std::mutex> lock(mRuntimeLayerMutex);
		std::string error;
		mRuntimeLayerModel.MarkBuildStarted(layerId, "Runtime Slang build started for shader '" + shaderId + "'.", error, preserveExistingRenderArtifact);
	}

	{
		std::lock_guard<std::mutex> lock(mShaderBuildMutex);
		mQueuedShaderBuilds.push_back(QueuedShaderBuild{ layerId, shaderId });
	}
	PumpShaderBuildQueue();
}

void RuntimeLayerController::PumpShaderBuildQueue()
{
	for (;;)
	{
		QueuedShaderBuild build;
		{
			std::lock_guard<std::mutex> lock(mShaderBuildMutex);
			if (mShaderBuilds.size() >= kMaxConcurrentRuntimeShaderBuilds || mQueuedShaderBuilds.empty())
				return;
			build = mQueuedShaderBuilds.front();
			mQueuedShaderBuilds.pop_front();
			if (mShaderBuilds.find(build.layerId) != mShaderBuilds.end())
				continue;
		}

		StartQueuedShaderBuild(build.layerId, build.shaderId);
	}
}

void RuntimeLayerController::StartQueuedShaderBuild(const std::string& layerId, const std::string& shaderId)
{
	auto bridge = std::make_unique<RuntimeShaderBridge>();
	RuntimeShaderBridge* bridgePtr = bridge.get();
	{
		std::lock_guard<std::mutex> lock(mShaderBuildMutex);
		if (mShaderBuilds.size() >= kMaxConcurrentRuntimeShaderBuilds || mShaderBuilds.find(layerId) != mShaderBuilds.end())
		{
			mQueuedShaderBuilds.push_front(QueuedShaderBuild{ layerId, shaderId });
			return;
		}
		mShaderBuilds[layerId] = std::move(bridge);
	}

	bridgePtr->Start(
		layerId,
		shaderId,
		[this, layerId](const RuntimeShaderArtifact& artifact) {
			if (MarkRuntimeBuildReady(artifact))
				PublishRuntimeRenderLayers();
			CompleteLayerShaderBuild(layerId);
		},
		[this, layerId](const std::string& message) {
			MarkRuntimeBuildFailedForLayer(layerId, message);
			LogError("runtime-shader", "Runtime Slang build failed: " + message);
			CompleteLayerShaderBuild(layerId);
		});
}

void RuntimeLayerController::CompleteLayerShaderBuild(const std::string& layerId)
{
	std::unique_ptr<RuntimeShaderBridge> bridge;
	{
		std::lock_guard<std::mutex> lock(mShaderBuildMutex);
		auto bridgeIt = mShaderBuilds.find(layerId);
		if (bridgeIt != mShaderBuilds.end())
		{
			bridge = std::move(bridgeIt->second);
			mShaderBuilds.erase(bridgeIt);
			mRetiredShaderBuilds.push_back(std::move(bridge));
		}
	}
	PumpShaderBuildQueue();
}

void RuntimeLayerController::RetireLayerShaderBuild(const std::string& layerId)
{
	std::unique_ptr<RuntimeShaderBridge> bridge;
	{
		std::lock_guard<std::mutex> lock(mShaderBuildMutex);
		mQueuedShaderBuilds.erase(
			std::remove_if(
				mQueuedShaderBuilds.begin(),
				mQueuedShaderBuilds.end(),
				[&layerId](const QueuedShaderBuild& build) {
					return build.layerId == layerId;
				}),
			mQueuedShaderBuilds.end());

		auto bridgeIt = mShaderBuilds.find(layerId);
		if (bridgeIt == mShaderBuilds.end())
			return;
		bridge = std::move(bridgeIt->second);
		mShaderBuilds.erase(bridgeIt);
		bridge->RequestStop();
		mRetiredShaderBuilds.push_back(std::move(bridge));
	}
}

void RuntimeLayerController::CleanupRetiredShaderBuilds()
{
	std::vector<std::unique_ptr<RuntimeShaderBridge>> readyToStop;
	{
		std::lock_guard<std::mutex> lock(mShaderBuildMutex);
		for (auto it = mRetiredShaderBuilds.begin(); it != mRetiredShaderBuilds.end();)
		{
			if ((*it)->CanStopWithoutWaiting())
			{
				readyToStop.push_back(std::move(*it));
				it = mRetiredShaderBuilds.erase(it);
				continue;
			}
			++it;
		}
	}

	for (std::unique_ptr<RuntimeShaderBridge>& bridge : readyToStop)
		bridge->Stop();
}

void RuntimeLayerController::StopAllRuntimeShaderBuilds()
{
	std::map<std::string, std::unique_ptr<RuntimeShaderBridge>> builds;
	std::vector<std::unique_ptr<RuntimeShaderBridge>> retiredBuilds;
	{
		std::lock_guard<std::mutex> lock(mShaderBuildMutex);
		mQueuedShaderBuilds.clear();
		builds.swap(mShaderBuilds);
		retiredBuilds.swap(mRetiredShaderBuilds);
	}
	for (auto& entry : builds)
		entry.second->Stop();
	for (auto& bridge : retiredBuilds)
		bridge->Stop();
}
}
