#include "RuntimeLayerController.h"

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

#include <filesystem>
#include <fstream>
#include <sstream>

namespace RenderCadenceCompositor
{
bool RuntimeLayerController::LoadSupportedShaderCatalog(const std::string& shaderLibrary, unsigned maxTemporalHistoryFrames)
{
	const std::filesystem::path shaderRoot = FindRepoPath(shaderLibrary);
	std::string error;
	if (!mShaderCatalog.Load(shaderRoot, maxTemporalHistoryFrames, error))
	{
		LogWarning("runtime-shader", "Supported shader catalog is empty: " + error);
		return false;
	}

	std::size_t preparedFontAtlases = 0;
	for (const auto& entry : mShaderCatalog.FontAtlases())
		preparedFontAtlases += entry.second.size();

	Log(
		"runtime-shader",
		"Supported shader catalog loaded with " + std::to_string(mShaderCatalog.Shaders().size()) +
			" shader(s), prepared " + std::to_string(preparedFontAtlases) + " font atlas asset(s).");
	return true;
}

void RuntimeLayerController::InitializeLayerModel(std::string& runtimeShaderId)
{
	if (InitializeLayerModelFromRuntimeState())
		return;

	if (!runtimeShaderId.empty())
		Log("runtime-state", "Falling back to configured runtime shader '" + runtimeShaderId + "'.");

	std::lock_guard<std::mutex> lock(mRuntimeLayerMutex);
	std::string error;
	if (!mRuntimeLayerModel.InitializeSingleLayer(mShaderCatalog, runtimeShaderId, error))
	{
		LogWarning("runtime-shader", error + " Runtime shader build disabled.");
		runtimeShaderId.clear();
		mRuntimeLayerModel.Clear();
	}
}

bool RuntimeLayerController::InitializeLayerModelFromRuntimeState()
{
	const std::filesystem::path runtimeStatePath = mRuntimeStatePath.empty() ? ResolveRuntimeStatePath() : mRuntimeStatePath;
	if (runtimeStatePath.empty())
		return false;
	if (!std::filesystem::exists(runtimeStatePath))
		return false;

	std::ifstream input(runtimeStatePath, std::ios::binary);
	if (!input)
	{
		LogWarning("runtime-state", "Could not open runtime state file: " + runtimeStatePath.string());
		return false;
	}

	std::ostringstream buffer;
	buffer << input.rdbuf();

	JsonValue runtimeState;
	std::string error;
	if (!ParseJson(buffer.str(), runtimeState, error))
	{
		LogWarning("runtime-state", "Could not parse runtime state file: " + error);
		return false;
	}

	{
		std::lock_guard<std::mutex> lock(mRuntimeLayerMutex);
		if (!mRuntimeLayerModel.InitializeFromRuntimeState(mShaderCatalog, runtimeState, error))
		{
			LogWarning("runtime-state", "Could not restore runtime state: " + error);
			return false;
		}
	}

	Log("runtime-state", "Restored runtime layer stack from " + runtimeStatePath.string() + ".");
	return true;
}

std::filesystem::path RuntimeLayerController::ResolveRuntimeStatePath() const
{
	const std::filesystem::path runtimeDirectory = FindRepoPath("runtime");
	if (!runtimeDirectory.empty())
		return runtimeDirectory / "runtime_state.json";
	return std::filesystem::current_path() / "runtime" / "runtime_state.json";
}

void RuntimeLayerController::RequestRuntimeStatePersistence()
{
	std::lock_guard<std::mutex> lock(mRuntimeLayerMutex);
	RequestRuntimeStatePersistenceLocked();
}

void RuntimeLayerController::RequestRuntimeStatePersistenceLocked()
{
	mPersistenceWriter.RequestSave(mRuntimeLayerModel.Snapshot());
}

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

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

	auto bridge = std::make_unique<RuntimeShaderBridge>();
	RuntimeShaderBridge* bridgePtr = bridge.get();
	{
		std::lock_guard<std::mutex> lock(mShaderBuildMutex);
		mShaderBuilds[layerId] = std::move(bridge);
	}

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

void RuntimeLayerController::RetireLayerShaderBuild(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())
			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);
		builds.swap(mShaderBuilds);
		retiredBuilds.swap(mRetiredShaderBuilds);
	}
	for (auto& entry : builds)
		entry.second->Stop();
	for (auto& bridge : retiredBuilds)
		bridge->Stop();
}

std::string RuntimeLayerController::FirstRuntimeLayerId() const
{
	std::lock_guard<std::mutex> lock(mRuntimeLayerMutex);
	return mRuntimeLayerModel.FirstLayerId();
}
}