/*
===========================================================================
Copyright (C) 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 - full opaque pre-pass


#include "crp_local.h"
#include "compshaders/crp/prepass.h"


#pragma pack(push, 4)
struct PrepassRC
{
	float modelViewMatrix[16];
	float modelMatrix[16];
	float normalMatrix[16];
	uint32_t textureIndex;
	uint32_t samplerIndex;
	uint32_t alphaTest;
};
#pragma pack(pop)


void Prepass::Init()
{
	psoCache.Init(psoCacheEntries, ARRAY_LEN(psoCacheEntries));
}

void Prepass::Draw(const drawSceneViewCommand_t& cmd)
{
	if(cmd.numDrawSurfs - cmd.numTranspSurfs <= 0)
	{
		return;
	}

	srp.renderMode = RenderMode::World;

	SCOPED_RENDER_PASS("Pre-pass", 1.0f, 0.5f, 0.5f);

	backEnd.refdef = cmd.refdef;
	backEnd.viewParms = cmd.viewParms;

	CmdSetViewportAndScissor(backEnd.viewParms);
	batchOldDepthHack = false;
	batchDepthHack = false;

	CmdBeginBarrier();
	CmdTextureBarrier(crp.depthTexture, ResourceStates::DepthWriteBit);
	CmdTextureBarrier(crp.normalTexture, ResourceStates::RenderTargetBit);
	CmdTextureBarrier(crp.motionVectorTexture, ResourceStates::RenderTargetBit);
	CmdTextureBarrier(crp.shadingPositionTexture, ResourceStates::RenderTargetBit);
	CmdEndBarrier();

	CmdClearDepthStencilTarget(crp.depthTexture, true, 0.0f);
	CmdClearColorTarget(crp.normalTexture, vec4_zero, NULL);
	CmdClearColorTarget(crp.motionVectorTexture, vec4_zero, NULL);
	CmdClearColorTarget(crp.shadingPositionTexture, vec4_zero, NULL);

	GeoBuffers& db = crp.dynBuffers[GetFrameIndex()];
	db.BeginUpload();

	const HTexture renderTargets[] = { crp.normalTexture, crp.motionVectorTexture, crp.shadingPositionTexture };
	CmdBindRenderTargets(ARRAY_LEN(renderTargets), renderTargets, &crp.depthTexture);
	CmdBindVertexBuffers(ARRAY_LEN(db.vertexBuffers), db.vertexBuffers, db.vertexBufferStrides, NULL);
	CmdBindIndexBuffer(db.indexBuffer.buffer, IndexType::UInt32, 0);

	const drawSurf_t* drawSurfs = cmd.drawSurfs;
	const int surfCount = cmd.numDrawSurfs - cmd.numTranspSurfs;
	const double originalTime = backEnd.refdef.floatTime;

	const shader_t* shader = NULL;
	const shader_t* oldShader = NULL;
	int oldEntityNum = -1;
	backEnd.currentEntity = &tr.worldEntity;

	tess.numVertexes = 0;
	tess.numIndexes = 0;

	int ds;
	const drawSurf_t* drawSurf;
	for(ds = 0, drawSurf = drawSurfs; ds < surfCount; ++ds, ++drawSurf)
	{
		int entityNum;
		R_DecomposeSort(drawSurf->sort, &entityNum, &shader);
		Q_assert(shader != NULL);
		Q_assert(shader->isOpaque);

		if(shader->isSky || shader->numStages <= 0)
		{
			continue;
		}

		const bool shaderChanged = shader != oldShader;
		const bool entityChanged = entityNum != oldEntityNum;
		if(shaderChanged || entityChanged)
		{
			oldShader = shader;
			oldEntityNum = entityNum;
			EndBatch();
			BeginBatch(shader);
			tess.greyscale = drawSurf->greyscale;
		}

		if(entityChanged)
		{
			UpdateEntityData(batchDepthHack, entityNum, originalTime);
		}

		R_TessellateSurface(drawSurf->surface);
	}

	backEnd.refdef.floatTime = originalTime;

	EndBatch();

	db.EndUpload();

	// restores the potentially "hacked" depth range as well
	CmdSetViewportAndScissor(backEnd.viewParms);
	batchOldDepthHack = false;
	batchDepthHack = false;
}

void Prepass::ProcessShader(shader_t& shader)
{
	if(shader.isSky)
	{
		return;
	}

	Q_assert(shader.isOpaque);

	if(shader.numStages < 1)
	{
		return;
	}

	const bool clampDepth = r_depthClamp->integer != 0 || shader.isSky;
	int a = 0;

#if defined(DEBUG)
	const shaderStage_t& stage = *shader.stages[0];
	const unsigned int stateBits = stage.stateBits & (~GLS_POLYMODE_LINE);
	Q_assert((stateBits & GLS_DEPTHTEST_DISABLE) == 0); // depth test enabled
	Q_assert((stateBits & GLS_DEPTHMASK_TRUE) != 0); // depth write enabled
	Q_assert((stateBits & GLS_DEPTHFUNC_EQUAL) == 0); // depth comparison GE
#endif

	// @NOTE: we 0-init the struct so that padding bytes don't mess up comparisons in the PSO cache
	GraphicsPipelineDesc desc = {};
	desc.name = "pre-pass";
	desc.rootSignature = RHI_MAKE_NULL_HANDLE();
	desc.shortLifeTime = true; // the PSO cache is only valid for this map!
	desc.vertexShader.Set(g_prepass_vs);
	desc.pixelShader.Set(g_prepass_ps);
	desc.vertexLayout.AddAttribute(a++, ShaderSemantic::Position, DataType::Float32, 3, 0);
	desc.vertexLayout.AddAttribute(a++, ShaderSemantic::Normal, DataType::Float32, 3, 0);
	desc.vertexLayout.AddAttribute(a++, ShaderSemantic::TexCoord, DataType::Float32, 2, 0);
	desc.vertexLayout.AddAttribute(a++, ShaderSemantic::Color, DataType::UNorm8, 4, 0);
	desc.depthStencil.depthStencilFormat = TextureFormat::Depth32_Float;
	desc.depthStencil.depthComparison = ComparisonFunction::GreaterEqual;
	desc.depthStencil.enableDepthTest = true;
	desc.depthStencil.enableDepthWrites = true;
	desc.rasterizer.cullMode = shader.cullType;
	desc.rasterizer.polygonOffset = shader.polygonOffset != 0;
	desc.rasterizer.clampDepth = clampDepth;
	desc.AddRenderTarget(0, TextureFormat::RG32_SNorm);
	desc.AddRenderTarget(0, TextureFormat::RG32_Float);
	desc.AddRenderTarget(0, TextureFormat::RGBA128_Float);

	pipeline_t& p = shader.prepassPipeline;
	p.firstStage = 0;
	p.numStages = 1;
	p.pipeline = psoCache.AddPipeline(desc, va("pre-pass %d", psoCache.entryCount));
	desc.rasterizer.cullMode = GetMirrorredCullType(desc.rasterizer.cullMode);
	p.mirrorPipeline = psoCache.AddPipeline(desc, va("pre-pass %d mirrored", psoCache.entryCount));
}

void Prepass::TessellationOverflow()
{
	EndBatch();
	BeginBatch(tess.shader);
}

void Prepass::BeginBatch(const shader_t* shader)
{
	tess.tessellator = Tessellator::Prepass;
	tess.numVertexes = 0;
	tess.numIndexes = 0;
	tess.depthFade = DFT_NONE;
	tess.deformsPreApplied = qfalse;
	tess.xstages = (const shaderStage_t**)shader->stages;
	tess.shader = shader;
	tess.shaderTime = backEnd.refdef.floatTime - tess.shader->timeOffset;
	if(tess.shader->clampTime && tess.shaderTime >= tess.shader->clampTime)
	{
		tess.shaderTime = tess.shader->clampTime;
	}
}

void Prepass::EndBatch()
{
	PrepassRC rc = {};
	float tempMatrix[16];

	const int vertexCount = tess.numVertexes;
	const int indexCount = tess.numIndexes;
	if(vertexCount <= 0 ||
		indexCount <= 0 ||
		tess.shader->numStages <= 0)
	{
		goto clean_up;
	}

	const shader_t* const shader = tess.shader;

	GeoBuffers& db = crp.dynBuffers[GetFrameIndex()];
	if(!db.CanAdd(vertexCount, indexCount, shader->numStages))
	{
		Q_assert(!"World surface geometry buffer too small!");
		goto clean_up;
	}

	RB_DeformTessGeometry(0, vertexCount, 0, indexCount);
	db.UploadBase();

	if(batchDepthHack != batchOldDepthHack)
	{
		const viewParms_t& vp = backEnd.viewParms;
		CmdSetViewport(vp.viewportX, vp.viewportY, vp.viewportWidth, vp.viewportHeight, batchDepthHack ? 0.7f : 0.0f, 1.0f);
		batchOldDepthHack = batchDepthHack;
	}

	const shaderStage_t* const stage = shader->stages[0];

	R_ComputeColors(stage, tess.svars[0], 0, vertexCount);
	R_ComputeTexCoords(stage, tess.svars[0], 0, vertexCount, qfalse);
	db.UploadStage(0);

	const pipeline_t& pipeline = shader->prepassPipeline;
	const int psoIndex = backEnd.viewParms.isMirror ? pipeline.mirrorPipeline : pipeline.pipeline;
	Q_assert(psoIndex > 0);
	CmdBindPipeline(psoCache.entries[psoIndex].handle);

	const image_t* image = GetBundleImage(stage->bundle);
	const uint32_t texIdx = image->textureIndex;
	const uint32_t sampIdx = GetSamplerIndex(image);
	const uint32_t alphaTest = AlphaTestShaderConstFromStateBits(stage->stateBits);
	Q_assert(sampIdx < ARRAY_LEN(crp.samplers));

	memcpy(rc.modelViewMatrix, backEnd.orient.modelMatrix, sizeof(rc.modelViewMatrix));
	memcpy(rc.modelMatrix, backEnd.modelMatrix, sizeof(rc.modelMatrix));
	R_InvMatrix(backEnd.modelMatrix, tempMatrix);
	R_TransposeMatrix(tempMatrix, rc.normalMatrix);
	rc.textureIndex = texIdx;
	rc.samplerIndex = sampIdx;
	rc.alphaTest = alphaTest;
	CmdSetGraphicsRootConstants(0, sizeof(rc), &rc);

	db.DrawStage(vertexCount, indexCount);

	db.EndBaseBatch(vertexCount);

clean_up:
	tess.tessellator = Tessellator::None;
	tess.numVertexes = 0;
	tess.numIndexes = 0;
}