/*
===========================================================================
Copyright (C) 2023-2024 Gian 'myT' Schellenbaum

This file is part of Challenge Quake 3 (CNQ3).

Challenge Quake 3 is free software; you can redistribute it
and/or modify it under the terms of the GNU General Public License as
published by the Free Software Foundation; either version 2 of the License,
or (at your option) any later version.

Challenge Quake 3 is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with Challenge Quake 3. If not, see <https://www.gnu.org/licenses/>.
===========================================================================
*/
// Cinematic Rendering Pipeline - scatter-as-gather depth of field


#include "crp_local.h"
namespace debug
{
#include "compshaders/crp/gatherdof_debug_vs.h"
#include "compshaders/crp/gatherdof_debug_ps.h"
}
namespace split
{
#include "compshaders/crp/gatherdof_split.h"
}
namespace near_coc_tile_gen
{
#include "compshaders/crp/gatherdof_coc_tile_gen.h"
}
namespace near_coc_tile_max
{
#include "compshaders/crp/gatherdof_coc_tile_max.h"
}
namespace blur
{
#include "compshaders/crp/gatherdof_blur.h"
}
namespace fill
{
#include "compshaders/crp/gatherdof_fill.h"
}
namespace combine
{
#include "compshaders/crp/gatherdof_combine_vs.h"
#include "compshaders/crp/gatherdof_combine_ps.h"
}


#pragma pack(push, 4)

struct DOFDebugRC
{
	uint32_t colorTextureIndex;
	uint32_t depthTextureIndex;
	uint32_t debugMode;
	float linearDepthA;
	float linearDepthB;
	float focusNearMin;
	float focusNearMax;
	float focusFarMin;
	float focusFarMax;
	float focusDist;
};

struct DOFSplitRC
{
	uint32_t depthTextureIndex;
	uint32_t colorTextureIndex;
	uint32_t nearColorTextureIndex;
	uint32_t farColorTextureIndex;
	uint32_t nearCocTextureIndex;
	uint32_t farCocTextureIndex;
	float linearDepthA;
	float linearDepthB;
	float focusNearMin;
	float focusNearMax;
	float focusFarMin;
	float focusFarMax;
	float brightnessScale;
};

struct DOFNearCocMaxRC
{
	uint32_t inputTextureIndex;
	uint32_t outputTextureIndex;
	uint32_t samplerIndex;
	int32_t kernelRadius;
	float kernelDirectionX;
	float kernelDirectionY;
};

struct DOFNearCocBlurRC
{
	uint32_t inputTextureIndex;
	uint32_t outputTextureIndex;
	uint32_t samplerIndex;
	int32_t kernelRadius;
	float kernelDirectionX;
	float kernelDirectionY;
};

struct DOFNearCocTileGenRC
{
	uint32_t inputTextureIndex;
	uint32_t outputTextureIndex;
};

struct DOFNearCocTileMaxRC
{
	uint32_t inputTextureIndex;
	uint32_t outputTextureIndex;
	uint32_t samplerIndex; // point/clamp
};

struct DOFBlurRC
{
	uint32_t colorTextureIndex;
	uint32_t nearColorTextureIndex;
	uint32_t nearMaxCocTextureIndex;
	uint32_t nearCocTextureIndex; // blurry
	uint32_t nearOutputTextureIndex;
	uint32_t farColorTextureIndex;
	uint32_t farCocTextureIndex; // sharp
	uint32_t farOutputTextureIndex;
	uint32_t samplerIndex; // linear/clamp
	float brightnessScale;
	float bladeCount;
	float bokehAngleRad;
};

struct DOFFillRC
{
	uint32_t nearInputTextureIndex;
	uint32_t nearOutputTextureIndex;
	uint32_t farInputTextureIndex;
	uint32_t farOutputTextureIndex;
	uint32_t samplerIndex; // point/clamp
};

struct DOFCombineRC
{
	uint32_t nearTextureIndex;
	uint32_t farTextureIndex;
	uint32_t nearCocTextureIndex;
	uint32_t farCocTextureIndex;
	uint32_t sharpTextureIndex;
	uint32_t samplerIndex; // point/clamp
};

#pragma pack(pop)


void GatherDepthOfField::Init()
{
	const TextureFormat::Id renderTargetFormat = TextureFormat::RGBA64_Float;

	tileWidth = (uint32_t)(glConfig.vidWidth + 15) / 16;
	tileHeight = (uint32_t)(glConfig.vidHeight + 15) / 16;

	{
		ComputePipelineDesc desc("DOF split");
		desc.shortLifeTime = true;
		desc.shader = ShaderByteCode(split::g_cs);
		splitPipeline = CreateComputePipeline(desc);
	}

	{
		ComputePipelineDesc desc("DOF near CoC tile generation");
		desc.shortLifeTime = true;
		desc.shader = ShaderByteCode(near_coc_tile_gen::g_cs);
		nearCocTileGenPipeline = CreateComputePipeline(desc);
	}

	{
		ComputePipelineDesc desc("DOF near CoC tile dilation");
		desc.shortLifeTime = true;
		desc.shader = ShaderByteCode(near_coc_tile_max::g_cs);
		nearCocTileMaxPipeline = CreateComputePipeline(desc);
	}

	{
		ComputePipelineDesc desc("DOF blur");
		desc.shortLifeTime = true;
		desc.shader = ShaderByteCode(blur::g_cs);
		blurPipeline = CreateComputePipeline(desc);
	}

	{
		ComputePipelineDesc desc("DOF fill");
		desc.shortLifeTime = true;
		desc.shader = ShaderByteCode(fill::g_cs);
		fillPipeline = CreateComputePipeline(desc);
	}

	{
		GraphicsPipelineDesc desc("DOF combine");
		desc.shortLifeTime = true;
		desc.vertexShader = ShaderByteCode(combine::g_vs);
		desc.pixelShader = ShaderByteCode(combine::g_ps);
		desc.depthStencil.DisableDepth();
		desc.rasterizer.cullMode = CT_TWO_SIDED;
		desc.AddRenderTarget(0, renderTargetFormat);
		combinePipeline = CreateGraphicsPipeline(desc);
	}

	{
		GraphicsPipelineDesc desc("DOF viz");
		desc.shortLifeTime = true;
		desc.vertexShader = ShaderByteCode(debug::g_vs);
		desc.pixelShader = ShaderByteCode(debug::g_ps);
		desc.depthStencil.DisableDepth();
		desc.rasterizer.cullMode = CT_TWO_SIDED;
		desc.AddRenderTarget(0, renderTargetFormat);
		debugPipeline = CreateGraphicsPipeline(desc);
	}

	{
		TextureDesc desc("DOF far field color", glConfig.vidWidth, glConfig.vidHeight);
		desc.shortLifeTime = true;
		desc.committedResource = true;
		desc.initialState = ResourceStates::UnorderedAccessBit;
		desc.allowedState = ResourceStates::UnorderedAccessBit | ResourceStates::ComputeShaderAccessBit | ResourceStates::PixelShaderAccessBit;
		desc.format = renderTargetFormat;
		farColorTexture = CreateTexture(desc);
		desc.name = "DOF near field color";
		nearColorTexture = CreateTexture(desc);
		desc.name = "DOF near field blurred color";
		nearBlurTexture = CreateTexture(desc);
		desc.name = "DOF far field blurred color";
		farBlurTexture = CreateTexture(desc);

		desc.format = TextureFormat::R8_UNorm;
		desc.name = "DOF near field CoC #1";
		nearCocTexture = CreateTexture(desc);
		desc.name = "DOF near field CoC #2";
		nearCocTexture2 = CreateTexture(desc);
		desc.name = "DOF far field CoC";
		farCocTexture = CreateTexture(desc);

		desc.width = tileWidth;
		desc.height = tileHeight;
		desc.name = "DOF near field CoC tile #1";
		nearCocTileTexture = CreateTexture(desc);
		desc.name = "DOF near field CoC tile #2";
		nearCocTileTexture2 = CreateTexture(desc);
	}
}

void GatherDepthOfField::Draw()
{
	if(crp_dof->integer != DOFMethod::Gather)
	{
		return;
	}

	if(backEnd.viewParms.viewportX != 0 ||
		backEnd.viewParms.viewportY != 0 ||
		backEnd.viewParms.viewportWidth != glConfig.vidWidth ||
		backEnd.viewParms.viewportHeight != glConfig.vidHeight)
	{
		return;
	}

	DrawSplit();
	DrawNearCocTileGen();
	DrawNearCocTileMax();
	DrawBlur();
	DrawFill();
	DrawCombine();
	DrawDebug();
}

void GatherDepthOfField::DrawDebug()
{
	if(crp_dof_overlay->integer == 0)
	{
		return;
	}

	SCOPED_RENDER_PASS("DOF Debug", 0.125f, 0.125f, 0.25f);

	crp.SwapRenderTargets();
	const TextureBarrier barriers[] =
	{
		TextureBarrier(crp.GetReadRenderTarget(), ResourceStates::PixelShaderAccessBit),
		TextureBarrier(crp.depthTexture, ResourceStates::PixelShaderAccessBit),
		TextureBarrier(crp.GetWriteRenderTarget(), ResourceStates::RenderTargetBit)
	};
	CmdBarrier(ARRAY_LEN(barriers), barriers);

	DOFDebugRC rc = {};
	rc.colorTextureIndex = GetTextureIndexSRV(crp.GetReadRenderTarget());
	rc.depthTextureIndex = GetTextureIndexSRV(crp.depthTexture);
	rc.debugMode = crp_dof_overlay->integer;
	RB_LinearDepthConstants(&rc.linearDepthA, &rc.linearDepthB);
	rc.focusNearMin = crp_gatherDof_focusNearDist->value - 0.5f * crp_gatherDof_focusNearRange->value;
	rc.focusNearMax = crp_gatherDof_focusNearDist->value + 0.5f * crp_gatherDof_focusNearRange->value;
	rc.focusFarMin = crp_gatherDof_focusFarDist->value - 0.5f * crp_gatherDof_focusFarRange->value;
	rc.focusFarMax = crp_gatherDof_focusFarDist->value + 0.5f * crp_gatherDof_focusFarRange->value;
	rc.focusDist = 0.5f * (rc.focusNearMax + rc.focusFarMin);

	CmdBindRenderTargets(1, &crp.renderTarget, NULL);
	CmdBindPipeline(debugPipeline);
	CmdSetGraphicsRootConstants(0, sizeof(rc), &rc);
	CmdDraw(3, 0);
}

void GatherDepthOfField::DrawSplit()
{
	SCOPED_RENDER_PASS("DOF Split", 0.125f, 0.125f, 0.25f);

	const TextureBarrier barriers[] =
	{
		TextureBarrier(crp.depthTexture, ResourceStates::ComputeShaderAccessBit),
		TextureBarrier(crp.renderTarget, ResourceStates::ComputeShaderAccessBit),
		TextureBarrier(nearColorTexture, ResourceStates::UnorderedAccessBit),
		TextureBarrier(farColorTexture, ResourceStates::UnorderedAccessBit),
		TextureBarrier(nearCocTexture, ResourceStates::UnorderedAccessBit),
		TextureBarrier(farCocTexture, ResourceStates::UnorderedAccessBit)
	};
	CmdBarrier(ARRAY_LEN(barriers), barriers);

	DOFSplitRC rc = {};
	rc.depthTextureIndex = GetTextureIndexSRV(crp.depthTexture);
	rc.colorTextureIndex = GetTextureIndexSRV(crp.renderTarget);
	rc.nearColorTextureIndex = GetTextureIndexUAV(nearColorTexture, 0);
	rc.farColorTextureIndex = GetTextureIndexUAV(farColorTexture, 0);
	rc.nearCocTextureIndex = GetTextureIndexUAV(nearCocTexture, 0);
	rc.farCocTextureIndex = GetTextureIndexUAV(farCocTexture, 0);
	RB_LinearDepthConstants(&rc.linearDepthA, &rc.linearDepthB);
	rc.focusNearMin = crp_gatherDof_focusNearDist->value - 0.5f * crp_gatherDof_focusNearRange->value;
	rc.focusNearMax = crp_gatherDof_focusNearDist->value + 0.5f * crp_gatherDof_focusNearRange->value;
	rc.focusFarMin = crp_gatherDof_focusFarDist->value - 0.5f * crp_gatherDof_focusFarRange->value;
	rc.focusFarMax = crp_gatherDof_focusFarDist->value + 0.5f * crp_gatherDof_focusFarRange->value;
	rc.brightnessScale = crp_gatherDof_brightness->value;

	CmdBindPipeline(splitPipeline);
	CmdSetComputeRootConstants(0, sizeof(rc), &rc);
	CmdDispatch((glConfig.vidWidth + 7) / 8, (glConfig.vidHeight + 7) / 8, 1);
}

void GatherDepthOfField::DrawNearCocTileGen()
{
	SCOPED_RENDER_PASS("DOF Tile Gen", 0.125f, 0.125f, 0.25f);

	const TextureBarrier barriers[] =
	{
		TextureBarrier(nearCocTexture, ResourceStates::ComputeShaderAccessBit),
		TextureBarrier(nearCocTileTexture, ResourceStates::UnorderedAccessBit)
	};
	CmdBarrier(ARRAY_LEN(barriers), barriers);

	DOFNearCocTileGenRC rc = {};
	rc.inputTextureIndex = GetTextureIndexSRV(nearCocTexture);
	rc.outputTextureIndex = GetTextureIndexUAV(nearCocTileTexture, 0);

	CmdBindPipeline(nearCocTileGenPipeline);
	CmdSetComputeRootConstants(0, sizeof(rc), &rc);
	CmdDispatch((tileWidth + 7) / 8, (tileHeight + 7) / 8, 1);
}

void GatherDepthOfField::DrawNearCocTileMax()
{
	SCOPED_RENDER_PASS("DOF Tile Max", 0.125f, 0.125f, 0.25f);

	const TextureBarrier barriers[] =
	{
		TextureBarrier(nearCocTileTexture, ResourceStates::ComputeShaderAccessBit),
		TextureBarrier(nearCocTileTexture2, ResourceStates::UnorderedAccessBit)
	};
	CmdBarrier(ARRAY_LEN(barriers), barriers);

	DOFNearCocTileMaxRC rc = {};
	rc.inputTextureIndex = GetTextureIndexSRV(nearCocTileTexture);
	rc.outputTextureIndex = GetTextureIndexUAV(nearCocTileTexture2, 0);
	rc.samplerIndex = GetSamplerIndex(TW_CLAMP_TO_EDGE, TextureFilter::Point);

	CmdBindPipeline(nearCocTileMaxPipeline);
	CmdSetComputeRootConstants(0, sizeof(rc), &rc);
	CmdDispatch((tileWidth + 7) / 8, (tileHeight + 7) / 8, 1);
}

void GatherDepthOfField::DrawBlur()
{
	SCOPED_RENDER_PASS("DOF Blur", 0.125f, 0.125f, 0.25f);

	const TextureBarrier barriers[] =
	{
		TextureBarrier(crp.renderTarget, ResourceStates::ComputeShaderAccessBit),
		TextureBarrier(nearColorTexture, ResourceStates::ComputeShaderAccessBit),
		TextureBarrier(farColorTexture, ResourceStates::ComputeShaderAccessBit),
		TextureBarrier(nearCocTexture, ResourceStates::ComputeShaderAccessBit),
		TextureBarrier(nearCocTileTexture2, ResourceStates::ComputeShaderAccessBit),
		TextureBarrier(farCocTexture, ResourceStates::ComputeShaderAccessBit),
		TextureBarrier(nearBlurTexture, ResourceStates::UnorderedAccessBit),
		TextureBarrier(farBlurTexture, ResourceStates::UnorderedAccessBit)
	};
	CmdBarrier(ARRAY_LEN(barriers), barriers);

	DOFBlurRC rc = {};
	rc.colorTextureIndex = GetTextureIndexSRV(crp.renderTarget);
	rc.nearColorTextureIndex = GetTextureIndexSRV(nearColorTexture);
	rc.nearMaxCocTextureIndex = GetTextureIndexSRV(nearCocTileTexture2);
	rc.nearCocTextureIndex = GetTextureIndexSRV(nearCocTexture);
	rc.nearOutputTextureIndex = GetTextureIndexUAV(nearBlurTexture, 0);
	rc.farColorTextureIndex = GetTextureIndexSRV(farColorTexture);
	rc.farCocTextureIndex = GetTextureIndexSRV(farCocTexture);
	rc.farOutputTextureIndex = GetTextureIndexUAV(farBlurTexture, 0);
	rc.samplerIndex = GetSamplerIndex(TW_CLAMP_TO_EDGE, TextureFilter::Linear);
	rc.brightnessScale = crp_gatherDof_brightness->value;
	rc.bladeCount = crp_dof_blades->value;
	rc.bokehAngleRad = DEG2RAD(crp_dof_angle->value);

	CmdBindPipeline(blurPipeline);
	CmdSetComputeRootConstants(0, sizeof(rc), &rc);
	CmdDispatch((glConfig.vidWidth + 7) / 8, (glConfig.vidHeight + 7) / 8, 1);
}

void GatherDepthOfField::DrawFill()
{
	SCOPED_RENDER_PASS("DOF Fill", 0.125f, 0.125f, 0.25f);

	const TextureBarrier barriers[] =
	{
		TextureBarrier(nearBlurTexture, ResourceStates::ComputeShaderAccessBit),
		TextureBarrier(farBlurTexture, ResourceStates::ComputeShaderAccessBit),
		TextureBarrier(nearColorTexture, ResourceStates::UnorderedAccessBit),
		TextureBarrier(farColorTexture, ResourceStates::UnorderedAccessBit)
	};
	CmdBarrier(ARRAY_LEN(barriers), barriers);

	DOFFillRC rc = {};
	rc.nearInputTextureIndex = GetTextureIndexSRV(nearBlurTexture);
	rc.farInputTextureIndex = GetTextureIndexSRV(farBlurTexture);
	rc.nearOutputTextureIndex = GetTextureIndexUAV(nearColorTexture, 0);
	rc.farOutputTextureIndex = GetTextureIndexUAV(farColorTexture, 0);
	rc.samplerIndex = GetSamplerIndex(TW_CLAMP_TO_EDGE, TextureFilter::Point);

	CmdBindPipeline(fillPipeline);
	CmdSetComputeRootConstants(0, sizeof(rc), &rc);
	CmdDispatch((glConfig.vidWidth + 7) / 8, (glConfig.vidHeight + 7) / 8, 1);
}

void GatherDepthOfField::DrawCombine()
{
	SCOPED_RENDER_PASS("DOF Combine", 0.125f, 0.125f, 0.25f);

	const TextureBarrier barriers[] =
	{
		TextureBarrier(nearColorTexture, ResourceStates::PixelShaderAccessBit),
		TextureBarrier(farColorTexture, ResourceStates::PixelShaderAccessBit),
		TextureBarrier(nearCocTexture, ResourceStates::PixelShaderAccessBit),
		TextureBarrier(farCocTexture, ResourceStates::PixelShaderAccessBit),
		TextureBarrier(crp.GetReadRenderTarget(), ResourceStates::PixelShaderAccessBit),
		TextureBarrier(crp.GetWriteRenderTarget(), ResourceStates::RenderTargetBit)
	};
	CmdBarrier(ARRAY_LEN(barriers), barriers);

	DOFCombineRC rc = {};
	rc.nearTextureIndex = GetTextureIndexSRV(nearColorTexture);
	rc.farTextureIndex = GetTextureIndexSRV(farColorTexture);
	rc.nearCocTextureIndex = GetTextureIndexSRV(nearCocTexture);
	rc.farCocTextureIndex = GetTextureIndexSRV(farCocTexture);
	rc.sharpTextureIndex = GetTextureIndexSRV(crp.renderTarget);
	rc.samplerIndex = GetSamplerIndex(TW_CLAMP_TO_EDGE, TextureFilter::Point);

	CmdBindRenderTargets(1, &crp.renderTarget, NULL);
	CmdBindPipeline(combinePipeline);
	CmdSetGraphicsRootConstants(0, sizeof(rc), &rc);
	CmdDraw(3, 0);
}