/*
===========================================================================
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/>.
===========================================================================
*/
// Shared Rendering Pipeline - texture mip-map generation


#include "srp_local.h"


#pragma pack(push, 4)

struct StartConstants
{
	float gamma;
	uint32_t srcTexture;
	uint32_t dstTexture;
};

struct DownConstants
{
	float weights[4];
	int32_t maxSize[2];
	int32_t scale[2];
	int32_t offset[2];
	uint32_t clampMode; // 0 = repeat
	uint32_t srcMip;
	uint32_t dstMip;
	uint32_t srcTexture;
	uint32_t dstTexture;
};

struct EndConstants
{
	float blendColor[4];
	float intensity;
	float invGamma; // 1.0 / gamma
	uint32_t srcMip;
	uint32_t dstMip;
	uint32_t srcTexture;
	uint32_t dstTexture;
};

#pragma pack(pop)


void MipMapGenerator::Init(bool ddhi_, const ShaderByteCode& g2l, const ShaderByteCode& down, const ShaderByteCode& l2g)
{
	if(!srp.firstInit)
	{
		return;
	}

	ddhi = ddhi_;

	for(int t = 0; t < 2; ++t)
	{
		TextureDesc desc(va("mip-map generation #%d", t + 1), MAX_TEXTURE_SIZE, MAX_TEXTURE_SIZE);
		desc.format = TextureFormat::RGBA64_Float;
		desc.initialState = ResourceStates::UnorderedAccessBit;
		desc.allowedState = ResourceStates::UnorderedAccessBit | ResourceStates::ComputeShaderAccessBit;
		textures[MipSlice::Float16_0 + t] = CreateTexture(desc);
	}

	const char* const stageNames[] = { "start", "down", "end" };
	const uint32_t stageRCByteCount[] = { sizeof(StartConstants), sizeof(DownConstants), sizeof(EndConstants) };
	const uint32_t stageExtraMips[] = { 1, 0, MaxTextureMips };
	const ShaderByteCode stageShaders[] = { g2l, down, l2g };

	for(int s = 0; s < 3; ++s)
	{
		Stage& stage = stages[s];
		if(ddhi)
		{
			stage.rootSignature = RHI_MAKE_NULL_HANDLE();
		}
		else
		{
			{
				RootSignatureDesc desc(va("mip-map %s", stageNames[s]));
				desc.pipelineType = PipelineType::Compute;
				desc.constants[ShaderStage::Compute].byteCount = stageRCByteCount[s];
				desc.AddRange(DescriptorType::RWTexture, 0, MipSlice::Count + stageExtraMips[s]);
				stage.rootSignature = CreateRootSignature(desc);
			}
			{
				const DescriptorTableDesc desc(DescriptorTableDesc(va("mip-map %s", stageNames[s]), stage.rootSignature));
				stage.descriptorTable = CreateDescriptorTable(desc);

				DescriptorTableUpdate update;
				update.SetRWTexturesSlice(ARRAY_LEN(textures), textures, 0, 0);
				UpdateDescriptorTable(stage.descriptorTable, update);
			}
		}
		{
			ComputePipelineDesc desc(va("mip-map %s", stageNames[s]), stage.rootSignature);
			desc.shader = stageShaders[s];
			stage.pipeline = CreateComputePipeline(desc);
		}
	}
}

void MipMapGenerator::GenerateMipMaps(HTexture texture)
{
	// @TODO: better look-up
	image_t* image = NULL;
	for(int i = 0; i < tr.numImages; ++i)
	{
		if(tr.images[i]->texture == texture)
		{
			image = tr.images[i];
			break;
		}
	}
	Q_assert(image);
	if(image == NULL || (image->flags & IMG_NOMIPMAP) != 0)
	{
		return;
	}

	const int mipCount = R_ComputeMipCount(image->width, image->height);
	if(mipCount <= 1)
	{
		return;
	}

	BeginTempCommandList();

	if(!ddhi)
	{
		// this must happen after the BeginTempCommandList call because
		// it has a CPU wait that guarantees it's safe to update the descriptor tables
		{
			Stage& stage = stages[Stage::Start];
			DescriptorTableUpdate update;
			update.SetRWTexturesSlice(1, &texture, MipSlice::Count, 0);
			UpdateDescriptorTable(stage.descriptorTable, update);
		}
		{
			Stage& stage = stages[Stage::End];
			DescriptorTableUpdate update;
			update.SetRWTexturesChain(1, &texture, MipSlice::Count);
			UpdateDescriptorTable(stage.descriptorTable, update);
		}
	}

	int w = image->width;
	int h = image->height;
	enum { GroupSize = 8, GroupMask = GroupSize - 1 };

#define Dispatch(Width, Height) CmdDispatch((Width + GroupMask) / GroupSize, (Height + GroupMask) / GroupSize, 1)

	// create a linear-space copy of mip 0 into float16 texture 0
	{
		Stage& stage = stages[Stage::Start];
		StartConstants rc = {};
		rc.gamma = r_mipGenGamma->value;

		CmdBeginBarrier();
		CmdTextureBarrier(textures[MipSlice::Float16_0], ResourceStates::UnorderedAccessBit);
		CmdTextureBarrier(texture, ResourceStates::UnorderedAccessBit);
		CmdEndBarrier();

		if(!ddhi)
		{
			CmdBindRootSignature(stage.rootSignature);
		}
		CmdBindPipeline(stage.pipeline);
		if(ddhi)
		{
			rc.srcTexture = GetTextureIndexUAV(texture, 0);
			rc.dstTexture = GetTextureIndexUAV(textures[MipSlice::Float16_0], 0);
			CmdSetComputeRootConstants(0, sizeof(rc), &rc);
		}
		else
		{
			CmdBindDescriptorTable(stage.rootSignature, stage.descriptorTable);
			CmdSetRootConstants(stage.rootSignature, ShaderStage::Compute, &rc);
		}
		Dispatch(w, h);
	}

	// overwrite mip 0 to apply r_intensity if needed
	if((image->flags & IMG_NOIMANIP) == 0 &&
		r_intensity->value != 1.0f)
	{
		Stage& stage = stages[Stage::End];
		const int destMip = 0;

		EndConstants rc = {};
		rc.intensity = r_intensity->value;
		rc.invGamma = 1.0f / r_mipGenGamma->value;
		Vector4Clear(rc.blendColor);
		rc.srcMip = MipSlice::Float16_0;
		rc.dstMip = MipSlice::Count + destMip;

		CmdBeginBarrier();
		CmdTextureBarrier(textures[MipSlice::Float16_0], ResourceStates::UnorderedAccessBit);
		CmdTextureBarrier(texture, ResourceStates::UnorderedAccessBit);
		CmdEndBarrier();

		if(!ddhi)
		{
			CmdBindRootSignature(stage.rootSignature);
		}
		CmdBindPipeline(stage.pipeline);
		if(ddhi)
		{
			rc.srcTexture = GetTextureIndexUAV(textures[MipSlice::Float16_0], 0);
			rc.dstTexture = GetTextureIndexUAV(texture, destMip);
			CmdSetComputeRootConstants(0, sizeof(rc), &rc);
		}
		else
		{
			CmdBindDescriptorTable(stage.rootSignature, stage.descriptorTable);
			CmdSetRootConstants(stage.rootSignature, ShaderStage::Compute, &rc);
		}
		Dispatch(w, h);
	}

	for(int i = 1; i < mipCount; ++i)
	{
		const int ws = w;
		const int hs = h;
		w = max(w / 2, 1);
		h = max(h / 2, 1);

		// down-sample the image into float16 texture 1 and then 0
		{
			Stage& stage = stages[Stage::DownSample];
			DownConstants rc = {};
			rc.clampMode = image->wrapClampMode == TW_REPEAT ? 0 : 1;
			memcpy(rc.weights, tr.mipFilter, sizeof(rc.weights));

			if(!ddhi)
			{
				CmdBindRootSignature(stage.rootSignature);
			}
			CmdBindPipeline(stage.pipeline);
			if(!ddhi)
			{
				CmdBindDescriptorTable(stage.rootSignature, stage.descriptorTable);
			}

			CmdBeginBarrier();
			CmdTextureBarrier(textures[MipSlice::Float16_0], ResourceStates::UnorderedAccessBit);
			CmdTextureBarrier(textures[MipSlice::Float16_1], ResourceStates::UnorderedAccessBit);
			CmdEndBarrier();

			// down-sample on the X-axis
			rc.srcMip = MipSlice::Float16_0;
			rc.dstMip = MipSlice::Float16_1;
			rc.scale[0] = ws / w;
			rc.scale[1] = 1;
			rc.maxSize[0] = ws - 1;
			rc.maxSize[1] = hs - 1;
			rc.offset[0] = 1;
			rc.offset[1] = 0;
			if(ddhi)
			{
				rc.srcTexture = GetTextureIndexUAV(textures[MipSlice::Float16_0], 0);
				rc.dstTexture = GetTextureIndexUAV(textures[MipSlice::Float16_1], 0);
				CmdSetComputeRootConstants(0, sizeof(rc), &rc);
			}
			else
			{
				CmdSetRootConstants(stage.rootSignature, ShaderStage::Compute, &rc);
			}
			Dispatch(w, hs);

			CmdBeginBarrier();
			CmdTextureBarrier(textures[MipSlice::Float16_0], ResourceStates::UnorderedAccessBit);
			CmdTextureBarrier(textures[MipSlice::Float16_1], ResourceStates::UnorderedAccessBit);
			CmdEndBarrier();

			// down-sample on the Y-axis
			rc.srcMip = MipSlice::Float16_1;
			rc.dstMip = MipSlice::Float16_0;
			rc.scale[0] = 1;
			rc.scale[1] = hs / h;
			rc.maxSize[0] = w - 1;
			rc.maxSize[1] = hs - 1;
			rc.offset[0] = 0;
			rc.offset[1] = 1;
			if(ddhi)
			{
				rc.srcTexture = GetTextureIndexUAV(textures[MipSlice::Float16_1], 0);
				rc.dstTexture = GetTextureIndexUAV(textures[MipSlice::Float16_0], 0);
				CmdSetComputeRootConstants(0, sizeof(rc), &rc);
			}
			else
			{
				CmdSetRootConstants(stage.rootSignature, ShaderStage::Compute, &rc);
			}
			Dispatch(w, h);
		}

		// save the results in gamma-space
		{
			Stage& stage = stages[Stage::End];
			const int destMip = i;

			EndConstants rc = {};
			rc.intensity = (image->flags & IMG_NOIMANIP) ? 1.0f : r_intensity->value;
			rc.invGamma = 1.0f / r_mipGenGamma->value;
			memcpy(rc.blendColor, r_mipBlendColors[r_colorMipLevels->integer ? destMip : 0], sizeof(rc.blendColor));
			rc.srcMip = MipSlice::Float16_0;
			rc.dstMip = MipSlice::Count + destMip;

			CmdBeginBarrier();
			CmdTextureBarrier(textures[MipSlice::Float16_0], ResourceStates::UnorderedAccessBit);
			CmdTextureBarrier(texture, ResourceStates::UnorderedAccessBit);
			CmdEndBarrier();

			if(!ddhi)
			{
				CmdBindRootSignature(stage.rootSignature);
			}
			CmdBindPipeline(stage.pipeline);
			if(ddhi)
			{
				rc.srcTexture = GetTextureIndexUAV(textures[MipSlice::Float16_0], 0);
				rc.dstTexture = GetTextureIndexUAV(texture, destMip);
				CmdSetComputeRootConstants(0, sizeof(rc), &rc);
			}
			else
			{
				CmdBindDescriptorTable(stage.rootSignature, stage.descriptorTable);
				CmdSetRootConstants(stage.rootSignature, ShaderStage::Compute, &rc);
			}
			Dispatch(w, h);
		}
	}

#undef Dispatch

	EndTempCommandList();
}