#include <stdio.h>
#include <SpirvInterpreter.h>

#define CHECK_RESULT(expr) do { \
	SpvResult check_result = (expr); \
	if (check_result != SPV_RESULT_SUCCESS) \
	{ \
		fprintf(stderr, "%s failed with %d\n", #expr, check_result); \
		return -1; \
	} \
} while (0)

static SpvResult ReadImageFloat4(SpvReadImageInfo info, SpvVec4f* dst)
{
	(void)info;
	(void)dst;
	return SPV_RESULT_UNSUPPORTED_SPIRV;
}

static SpvResult ReadImageInt4(SpvReadImageInfo info, SpvVec4u* dst)
{
	(void)info;
	(void)dst;
	return SPV_RESULT_UNSUPPORTED_SPIRV;
}

static SpvResult WriteImageFloat4(SpvWriteImageInfo info, SpvVec4f src)
{
	(void)info;
	(void)src;
	return SPV_RESULT_UNSUPPORTED_SPIRV;
}

static SpvResult WriteImageInt4(SpvWriteImageInfo info, SpvVec4u src)
{
	(void)info;
	(void)src;
	return SPV_RESULT_UNSUPPORTED_SPIRV;
}

static SpvResult SampleImageFloat4(SpvSampleImageInfo info, SpvVec4f* dst)
{
	(void)info;
	(void)dst;
	return SPV_RESULT_UNSUPPORTED_SPIRV;
}

static SpvResult SampleImageInt4(SpvSampleImageInfo info, SpvVec4u* dst)
{
	(void)info;
	(void)dst;
	return SPV_RESULT_UNSUPPORTED_SPIRV;
}

static SpvResult SampleImageDref(SpvSampleImageInfo info, float dref, float* dst)
{
	(void)info;
	(void)dref;
	(void)dst;
	return SPV_RESULT_UNSUPPORTED_SPIRV;
}

static SpvResult QueryImageSize(SpvQueryImageInfo info, SpvVec4u* dst)
{
	(void)info;
	(void)dst;
	return SPV_RESULT_UNSUPPORTED_SPIRV;
}

static SpvResult QueryImageLevels(void* driver_image, unsigned long* dst)
{
	(void)driver_image;
	(void)dst;
	return SPV_RESULT_UNSUPPORTED_SPIRV;
}

static SpvResult QueryImageSamples(void* driver_image, unsigned long* dst)
{
	(void)driver_image;
	(void)dst;
	return SPV_RESULT_UNSUPPORTED_SPIRV;
}

static SpvResult QueryImageLod(SpvQueryImageLodInfo info, SpvVec4f* dst)
{
	(void)info;
	(void)dst;
	return SPV_RESULT_UNSUPPORTED_SPIRV;
}

static const unsigned char shader_source[]  = {
	0x03, 0x02, 0x23, 0x07, 0x00, 0x00, 0x01, 0x00, 0x82, 0x10, 0x27, 0x00, 0x17, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x11, 0x00, 0x02, 0x00, 0x01, 0x00, 0x00, 0x00, 0x0e, 0x00, 0x03, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x0f, 0x00, 0x06, 0x00, 0x04, 0x00, 0x00, 0x00,
	0x0f, 0x00, 0x00, 0x00, 0x6d, 0x61, 0x69, 0x6e, 0x00, 0x00, 0x00, 0x00, 0x06, 0x00, 0x00, 0x00,
	0x10, 0x00, 0x03, 0x00, 0x0f, 0x00, 0x00, 0x00, 0x07, 0x00, 0x00, 0x00, 0x03, 0x00, 0x03, 0x00,
	0x09, 0x00, 0x00, 0x00, 0x00, 0x10, 0x40, 0x00, 0x04, 0x00, 0x05, 0x00, 0x56, 0x65, 0x72, 0x73,
	0x69, 0x6f, 0x6e, 0x3a, 0x20, 0x31, 0x2e, 0x31, 0x00, 0x00, 0x00, 0x00, 0x05, 0x00, 0x04, 0x00,
	0x07, 0x00, 0x00, 0x00, 0x4f, 0x75, 0x74, 0x70, 0x75, 0x74, 0x00, 0x00, 0x06, 0x00, 0x05, 0x00,
	0x07, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x63, 0x6f, 0x6c, 0x6f, 0x72, 0x00, 0x00, 0x00,
	0x05, 0x00, 0x04, 0x00, 0x06, 0x00, 0x00, 0x00, 0x63, 0x6f, 0x6c, 0x6f, 0x72, 0x00, 0x00, 0x00,
	0x05, 0x00, 0x04, 0x00, 0x0f, 0x00, 0x00, 0x00, 0x6d, 0x61, 0x69, 0x6e, 0x00, 0x00, 0x00, 0x00,
	0x47, 0x00, 0x04, 0x00, 0x06, 0x00, 0x00, 0x00, 0x1e, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x48, 0x00, 0x05, 0x00, 0x07, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x23, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x13, 0x00, 0x02, 0x00, 0x01, 0x00, 0x00, 0x00, 0x21, 0x00, 0x03, 0x00,
	0x02, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x16, 0x00, 0x03, 0x00, 0x03, 0x00, 0x00, 0x00,
	0x20, 0x00, 0x00, 0x00, 0x17, 0x00, 0x04, 0x00, 0x04, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00,
	0x04, 0x00, 0x00, 0x00, 0x20, 0x00, 0x04, 0x00, 0x05, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00,
	0x04, 0x00, 0x00, 0x00, 0x1e, 0x00, 0x03, 0x00, 0x07, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00,
	0x20, 0x00, 0x04, 0x00, 0x08, 0x00, 0x00, 0x00, 0x07, 0x00, 0x00, 0x00, 0x07, 0x00, 0x00, 0x00,
	0x15, 0x00, 0x04, 0x00, 0x09, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
	0x2b, 0x00, 0x04, 0x00, 0x09, 0x00, 0x00, 0x00, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x2b, 0x00, 0x04, 0x00, 0x03, 0x00, 0x00, 0x00, 0x0b, 0x00, 0x00, 0x00, 0x00, 0x00, 0x80, 0x3f,
	0x2b, 0x00, 0x04, 0x00, 0x03, 0x00, 0x00, 0x00, 0x0c, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x40,
	0x2b, 0x00, 0x04, 0x00, 0x03, 0x00, 0x00, 0x00, 0x0d, 0x00, 0x00, 0x00, 0x00, 0x00, 0x40, 0x40,
	0x2b, 0x00, 0x04, 0x00, 0x03, 0x00, 0x00, 0x00, 0x0e, 0x00, 0x00, 0x00, 0x00, 0x00, 0x80, 0x40,
	0x20, 0x00, 0x04, 0x00, 0x14, 0x00, 0x00, 0x00, 0x07, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00,
	0x3b, 0x00, 0x04, 0x00, 0x05, 0x00, 0x00, 0x00, 0x06, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00,
	0x36, 0x00, 0x05, 0x00, 0x01, 0x00, 0x00, 0x00, 0x0f, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x02, 0x00, 0x00, 0x00, 0xf8, 0x00, 0x02, 0x00, 0x10, 0x00, 0x00, 0x00, 0x3b, 0x00, 0x04, 0x00,
	0x08, 0x00, 0x00, 0x00, 0x11, 0x00, 0x00, 0x00, 0x07, 0x00, 0x00, 0x00, 0x50, 0x00, 0x07, 0x00,
	0x04, 0x00, 0x00, 0x00, 0x12, 0x00, 0x00, 0x00, 0x0b, 0x00, 0x00, 0x00, 0x0c, 0x00, 0x00, 0x00,
	0x0d, 0x00, 0x00, 0x00, 0x0e, 0x00, 0x00, 0x00, 0x41, 0x00, 0x05, 0x00, 0x14, 0x00, 0x00, 0x00,
	0x13, 0x00, 0x00, 0x00, 0x11, 0x00, 0x00, 0x00, 0x0a, 0x00, 0x00, 0x00, 0x3e, 0x00, 0x03, 0x00,
	0x13, 0x00, 0x00, 0x00, 0x12, 0x00, 0x00, 0x00, 0x3d, 0x00, 0x04, 0x00, 0x07, 0x00, 0x00, 0x00,
	0x15, 0x00, 0x00, 0x00, 0x11, 0x00, 0x00, 0x00, 0x51, 0x00, 0x05, 0x00, 0x04, 0x00, 0x00, 0x00,
	0x16, 0x00, 0x00, 0x00, 0x15, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x3e, 0x00, 0x03, 0x00,
	0x06, 0x00, 0x00, 0x00, 0x16, 0x00, 0x00, 0x00, 0xfd, 0x00, 0x01, 0x00, 0x38, 0x00, 0x01, 0x00
};

int main(void)
{
	SpvModule module;
	SpvModuleOptions options;
	options.use_simd_vectors_specializations = 1;

	if(SpvInitModule(&module, (SpvWord*)shader_source, sizeof(shader_source) / 4, options) != SPV_RESULT_SUCCESS)
	{
		fprintf(stderr, "Module init failed\n");
		return -1;
	}

	SpvWord binding_result = 0;
	if (SpvModuleGetBindingResult(module, 0, 0, &binding_result) != SPV_RESULT_NOT_FOUND)
	{
		fprintf(stderr, "Unexpected binding lookup result\n");
		return -1;
	}
	SpvModuleReflectionInfos reflection_infos = SpvModuleGetReflectionInfos(module);
	if (reflection_infos.has_atomics || reflection_infos.early_fragment_tests)
	{
		fprintf(stderr, "Unexpected reflection flags\n");
		return -1;
	}

	SpvImageAPI image_api = {
		.SpvReadImageFloat4 = ReadImageFloat4,
		.SpvReadImageInt4 = ReadImageInt4,
		.SpvWriteImageFloat4 = WriteImageFloat4,
		.SpvWriteImageInt4 = WriteImageInt4,
		.SpvSampleImageFloat4 = SampleImageFloat4,
		.SpvSampleImageInt4 = SampleImageInt4,
		.SpvSampleImageDref = SampleImageDref,
		.SpvQueryImageSize = QueryImageSize,
		.SpvQueryImageLevels = QueryImageLevels,
		.SpvQueryImageSamples = QueryImageSamples,
		.SpvQueryImageLod = QueryImageLod
	};

	SpvRuntime runtime;
	if(SpvInitRuntime(&runtime, module, image_api) != SPV_RESULT_SUCCESS)
	{
		fprintf(stderr, "Runtime init failed\n");
		return -1;
	}

	SpvRuntime second_runtime;
	if(SpvInitRuntime(&second_runtime, module, image_api) != SPV_RESULT_SUCCESS)
	{
		fprintf(stderr, "Second runtime init failed\n");
		SpvDeinitRuntime(runtime);
		return -1;
	}

	unsigned int spec_value = 64;
	SpvRuntimeSpecializationEntry spec_entry = {
		.id = 0,
		.offset = 0,
		.size = sizeof(spec_value)
	};
	CHECK_RESULT(SpvAddSpecializationInfo(runtime, spec_entry, (const SpvByte*)&spec_value, sizeof(spec_value)));
	CHECK_RESULT(SpvCopySpecializationConstantsFrom(second_runtime, runtime));

	SpvWord main_entry_index;
	CHECK_RESULT(SpvGetEntryPointByName(runtime, "main", &main_entry_index));
	SpvWord fragment_entry_index;
	CHECK_RESULT(SpvGetEntryPointByNameAndExecutionModel(runtime, "main", SpvExecutionModelFragment, &fragment_entry_index));
	if (fragment_entry_index != main_entry_index)
	{
		fprintf(stderr, "Unexpected execution-model entry point lookup\n");
		return -1;
	}
	CHECK_RESULT(SpvSelectEntryPoint(runtime, main_entry_index));
	CHECK_RESULT(SpvApplySpecializationLayout(runtime));
	CHECK_RESULT(SpvApplySpecializationInvocationLayout(second_runtime));
	SpvWorkgroupSize workgroup_size;
	CHECK_RESULT(SpvGetWorkgroupSize(runtime, &workgroup_size));
	if (workgroup_size.has_size)
	{
		fprintf(stderr, "Unexpected workgroup size\n");
		return -1;
	}
	SpvWorkgroupMemory workgroup_memory;
	CHECK_RESULT(SpvCreateWorkgroupMemory(runtime, &workgroup_memory));
	if (SpvGetWorkgroupMemoryCount(workgroup_memory) != 0)
	{
		fprintf(stderr, "Unexpected workgroup memory count\n");
		return -1;
	}
	SpvWorkgroupMemoryItem workgroup_memory_item;
	if (SpvGetWorkgroupMemoryItem(workgroup_memory, 0, &workgroup_memory_item) != SPV_RESULT_OUT_OF_BOUNDS)
	{
		fprintf(stderr, "Unexpected workgroup memory lookup result\n");
		return -1;
	}
	CHECK_RESULT(SpvBindWorkgroupMemory(runtime, workgroup_memory));
	SpvDestroyWorkgroupMemory(runtime, workgroup_memory);
	CHECK_RESULT(SpvCallEntryPoint(runtime, main_entry_index));

	float output[4];
	SpvWord output_result;
	CHECK_RESULT(SpvGetResultByName(runtime, "color", &output_result));
	CHECK_RESULT(SpvRefreshResultValueLayout(runtime, output_result));
	CHECK_RESULT(SpvReadOutput(runtime, (SpvByte*)output, sizeof(output), output_result));
	SpvWord builtin_result;
	if (SpvGetBuiltinResult(runtime, SpvBuiltInWorkgroupId, &builtin_result) != SPV_RESULT_NOT_FOUND)
	{
		fprintf(stderr, "Unexpected builtin lookup result\n");
		return -1;
	}

	SpvSize output_size = 0;
	CHECK_RESULT(SpvGetResultMemorySize(runtime, output_result, &output_size));
	if (output_size != sizeof(output))
	{
		fprintf(stderr, "Unexpected output metadata\n");
		SpvDeinitRuntime(second_runtime);
		SpvDeinitRuntime(runtime);
		SpvDeinitModule(module);
		return -1;
	}

	float dx[4] = { 1.0f, 2.0f, 3.0f, 4.0f };
	float dy[4] = { 5.0f, 6.0f, 7.0f, 8.0f };
	CHECK_RESULT(SpvSetDerivativeFromMemory(runtime, output_result, (const SpvByte*)dx, sizeof(dx), (const SpvByte*)dy, sizeof(dy)));
	CHECK_RESULT(SpvCopyDerivative(runtime, output_result, output_result));
	SpvClearDerivative(runtime, output_result);

	printf("Output: Vec4[%f, %f, %f, %f]\n", output[0], output[1], output[2], output[3]);

	SpvDeinitRuntime(second_runtime);
	SpvDeinitRuntime(runtime);
	SpvDeinitModule(module);
	return 0;
}