#include "RuntimeSlangShaderCompiler.h"

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

#include <chrono>
#include <filesystem>
#include <iostream>

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;
	}
}

bool IsStatelessSinglePassPackage(const ShaderPackage& shaderPackage, std::string& error)
{
	if (shaderPackage.passes.size() != 1)
	{
		error = "RenderCadenceCompositor currently supports only single-pass runtime shaders.";
		return false;
	}
	if (shaderPackage.temporal.enabled)
	{
		error = "RenderCadenceCompositor currently supports only stateless shaders; temporal history is not enabled in this app.";
		return false;
	}
	if (shaderPackage.feedback.enabled)
	{
		error = "RenderCadenceCompositor currently supports only stateless shaders; feedback storage is not enabled in this app.";
		return false;
	}
	if (!shaderPackage.textureAssets.empty())
	{
		error = "RenderCadenceCompositor does not load shader texture assets on the render thread; texture-backed shaders need a CPU-prepared asset handoff first.";
		return false;
	}
	if (!shaderPackage.fontAssets.empty())
	{
		error = "RenderCadenceCompositor does not load shader font assets on the render thread; text shaders need a CPU-prepared asset handoff first.";
		return false;
	}
	for (const ShaderParameterDefinition& parameter : shaderPackage.parameters)
	{
		if (parameter.type == ShaderParameterType::Text)
		{
			error = "RenderCadenceCompositor currently skips text parameters because they require per-shader text texture storage.";
			return false;
		}
	}
	return true;
}
}

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;
		}
		if (!IsStatelessSinglePassPackage(shaderPackage, error))
		{
			build.succeeded = false;
			build.message = error;
			return build;
		}

		const ShaderPassDefinition& pass = shaderPackage.passes.front();

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

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

		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;
		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;
	}
}