// Copyright (C) 2025 kanel
// This file is part of "Pulse"
// For conditions of distribution and use, see copyright notice in LICENSE

#include <Pulse.h>
#include "../../PulseInternal.h"
#include "D3D11.h"
#include "D3D11Device.h"
#include "D3D11ComputePipeline.h"

#include <d3dcompiler.h>
#include <string.h>

static HRESULT CompileComputeShader(PulseDevice device, const unsigned char* src, uint32_t src_size, const char* entry_point, ID3DBlob** blob)
{
	if(!src || !entry_point || !device || !blob)
	   return E_INVALIDARG;

	Direct3D11Device* d3d11_device = D3D11_RETRIEVE_DRIVER_DATA_AS(device, Direct3D11Device*);
	*blob = PULSE_NULLPTR;

	UINT flags = D3DCOMPILE_ENABLE_STRICTNESS;
	if(PULSE_IS_BACKEND_HIGH_LEVEL_DEBUG(device->backend))
		flags |= D3DCOMPILE_DEBUG;

	// We generally prefer to use the higher CS shader profile when possible as CS 5.0 is better performance on 11-class hardware
	LPCSTR profile = (ID3D11Device_GetFeatureLevel(d3d11_device->device) >= D3D_FEATURE_LEVEL_11_0) ? "cs_5_0" : "cs_4_0";

	ID3DBlob* shader_blob = PULSE_NULLPTR;
	ID3DBlob* error_blob = PULSE_NULLPTR;
	HRESULT hr = D3DCompile(src, src_size, PULSE_NULLPTR, PULSE_NULLPTR, PULSE_NULLPTR, entry_point, profile, flags, 0, &shader_blob, &error_blob);

	if(FAILED(hr))
	{
		if(error_blob)
		{
			if(PULSE_IS_BACKEND_LOW_LEVEL_DEBUG(device->backend))
				PulseLogInfoFmt(device->backend, "(D3D11) failed to compile HLSL shader. %s", ID3D10Blob_GetBufferPointer(error_blob));
			ID3D10Blob_Release(error_blob);
		}
		if(shader_blob)
			ID3D10Blob_Release(shader_blob);
		return hr;
	}

	*blob = shader_blob;
	return hr;
}

PulseComputePipeline Direct3D11CreateComputePipeline(PulseDevice device, const PulseComputePipelineCreateInfo* info)
{
	Direct3D11Device* d3d11_device = D3D11_RETRIEVE_DRIVER_DATA_AS(device, Direct3D11Device*);

	PulseComputePipelineHandler* pipeline = (PulseComputePipelineHandler*)calloc(1, sizeof(PulseComputePipelineHandler));
	PULSE_CHECK_ALLOCATION_RETVAL(pipeline, PULSE_NULL_HANDLE);

	Direct3D11ComputePipeline* d3d11_pipeline = (Direct3D11ComputePipeline*)calloc(1, sizeof(Direct3D11ComputePipeline));
	PULSE_CHECK_ALLOCATION_RETVAL(d3d11_pipeline, PULSE_NULL_HANDLE);

	pipeline->driver_data = d3d11_pipeline;

	ID3D10Blob* blob = PULSE_NULLPTR;
	if(info->format & PULSE_SHADER_FORMAT_HLSL_BIT)
		CHECK_D3D11_RETVAL(device->backend, CompileComputeShader(device, info->code, info->code_size, info->entrypoint, &blob), PULSE_ERROR_INITIALIZATION_FAILED, PULSE_NULL_HANDLE);

	if(blob == PULSE_NULLPTR)
	{
		D3DCreateBlob(info->code_size, &blob);
		memcpy(ID3D10Blob_GetBufferPointer(blob), (void*)info->code, info->code_size);
	}

	CHECK_D3D11_RETVAL(device->backend, ID3D11Device_CreateComputeShader(d3d11_device->device, ID3D10Blob_GetBufferPointer(blob), ID3D10Blob_GetBufferSize(blob), PULSE_NULLPTR, &d3d11_pipeline->shader), PULSE_ERROR_INITIALIZATION_FAILED, PULSE_NULL_HANDLE);;

	ID3D10Blob_Release(blob);

	if(PULSE_IS_BACKEND_HIGH_LEVEL_DEBUG(device->backend))
		PulseLogInfoFmt(device->backend, "(D3D11) created new compute pipeline %p", pipeline);
	return pipeline;
}

void Direct3D11DestroyComputePipeline(PulseDevice device, PulseComputePipeline pipeline)
{
	if(pipeline == PULSE_NULL_HANDLE)
	{
		if(PULSE_IS_BACKEND_LOW_LEVEL_DEBUG(device->backend))
			PulseLogWarning(device->backend, "compute pipeline is NULL, this may be a bug in your application");
		return;
	}
	Direct3D11ComputePipeline* d3d11_pipeline = D3D11_RETRIEVE_DRIVER_DATA_AS(pipeline, Direct3D11ComputePipeline*);
	ID3D11ComputeShader_Release(d3d11_pipeline->shader);
	free(d3d11_pipeline);
	if(PULSE_IS_BACKEND_HIGH_LEVEL_DEBUG(device->backend))
		PulseLogInfoFmt(device->backend, "(D3D11) destroyed compute pipeline %p", pipeline);
	free(pipeline);
}