#include "RuntimeSlangShaderCompiler.h"

#include "ShaderCompiler.h"
#include "ShaderPackageRegistry.h"
#include "ShaderTypes.h"
#include "SupportedShaderCatalog.h"

#include <chrono>
#include <filesystem>

namespace
{
std::filesystem::path FindRepoRoot()
{
	std::filesystem::path current = std::filesystem::current_path();
	for (;;)
	{
		if (std::filesystem::exists(current / "shaders" / "happy-accident" / "shader.slang") &&
			std::filesystem::exists(current / "runtime" / "templates" / "shader_wrapper.slang.in"))
		{
			return current;
		}

		const std::filesystem::path parent = current.parent_path();
		if (parent.empty() || parent == current)
			return std::filesystem::current_path();
		current = parent;
	}
}

}

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

void RuntimeSlangShaderCompiler::StartHappyAccidentBuild()
{
	StartShaderBuild("happy-accident");
}

void RuntimeSlangShaderCompiler::StartShaderBuild(const std::string& shaderId)
{
	if (mRunning.load(std::memory_order_acquire))
		return;

	if (mThread.joinable())
		mThread.join();

	{
		std::lock_guard<std::mutex> lock(mMutex);
		mReadyBuild = RuntimeSlangShaderBuild();
	}

	mRunning.store(true, std::memory_order_release);
	mThread = std::thread([this, shaderId]() {
		RuntimeSlangShaderBuild build = BuildShader(shaderId);
		{
			std::lock_guard<std::mutex> lock(mMutex);
			mReadyBuild = std::move(build);
			mReadyBuild.available = true;
		}
		mRunning.store(false, std::memory_order_release);
	});
}

void RuntimeSlangShaderCompiler::Stop()
{
	if (mThread.joinable())
		mThread.join();
	mRunning.store(false, std::memory_order_release);
}

bool RuntimeSlangShaderCompiler::TryConsume(RuntimeSlangShaderBuild& build)
{
	std::lock_guard<std::mutex> lock(mMutex);
	if (!mReadyBuild.available)
		return false;

	build = std::move(mReadyBuild);
	mReadyBuild = RuntimeSlangShaderBuild();
	return true;
}

RuntimeSlangShaderBuild RuntimeSlangShaderCompiler::BuildShader(const std::string& shaderId) const
{
	RuntimeSlangShaderBuild build;
	build.artifact.shaderId = shaderId;

	try
	{
		const std::filesystem::path repoRoot = FindRepoRoot();
		const std::filesystem::path shaderDir = repoRoot / "shaders" / shaderId;
		const std::filesystem::path runtimeBuildDir = repoRoot / "runtime" / "generated" / "render-cadence-compositor";

		ShaderPackageRegistry registry(0);
		ShaderPackage shaderPackage;
		std::string error;
		if (!registry.ParseManifest(shaderDir / "shader.json", shaderPackage, error))
		{
			build.succeeded = false;
			build.message = error.empty() ? "Shader manifest parse failed." : error;
			return build;
		}
		const RenderCadenceCompositor::ShaderSupportResult support =
			RenderCadenceCompositor::CheckStatelessSinglePassShaderSupport(shaderPackage);
		if (!support.supported)
		{
			build.succeeded = false;
			build.message = support.reason;
			return build;
		}

		ShaderCompiler compiler(
			repoRoot,
			runtimeBuildDir / (shaderId + ".wrapper.slang"),
			runtimeBuildDir / (shaderId + ".generated.glsl"),
			runtimeBuildDir / (shaderId + ".patched.glsl"),
			0);

		const auto start = std::chrono::steady_clock::now();
		for (const ShaderPassDefinition& pass : shaderPackage.passes)
		{
			std::string fragmentShaderSource;
			if (!compiler.BuildPassFragmentShaderSource(shaderPackage, pass, fragmentShaderSource, error))
			{
				build.succeeded = false;
				build.message = error.empty() ? "Slang compile failed." : error;
				return build;
			}

			RuntimeShaderPassArtifact passArtifact;
			passArtifact.passId = pass.id;
			passArtifact.fragmentShaderSource = std::move(fragmentShaderSource);
			passArtifact.inputNames = pass.inputNames;
			passArtifact.outputName = pass.outputName;
			build.artifact.passes.push_back(std::move(passArtifact));
		}

		const auto end = std::chrono::steady_clock::now();
		const double milliseconds = std::chrono::duration_cast<std::chrono::duration<double, std::milli>>(end - start).count();
		build.succeeded = true;
		build.artifact.shaderId = shaderPackage.id;
		build.artifact.displayName = shaderPackage.displayName;
		build.artifact.parameterDefinitions = shaderPackage.parameters;
		if (!build.artifact.passes.empty())
			build.artifact.fragmentShaderSource = build.artifact.passes.front().fragmentShaderSource;
		build.artifact.message = shaderPackage.displayName + " Slang compile completed in " + std::to_string(milliseconds) + " ms.";
		build.message = build.artifact.message;
		return build;
	}
	catch (const std::exception& exception)
	{
		build.succeeded = false;
		build.message = exception.what();
		return build;
	}
}