/*
===========================================================================
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"
#include "compshaders/crp/prepass_bary.h"
#include "compshaders/crp/fullscreen.h"
#include "compshaders/crp/skybox_motion.h"


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


struct EntityData
{
	float modelMatrix[16];
	int uniqueId;
};

struct EntityTracker
{
	void BeginFrame()
	{
		Q_assert(!begun);
		Q_assert(!enabled);

		const viewParms_t& vp = backEnd.viewParms;
		enabled = IsViewportFullscreen(vp) && !vp.isPortal;
		begun = true;
	}

	void EndFrame()
	{
		Q_assert(begun);

		begun = false;
		if(enabled)
		{
			enabled = false;
			EntityArray* prevEntsCopy = prevEnts;
			prevEnts = currEnts;
			currEnts = prevEntsCopy;
			currEnts->Clear();
		}
	}

	void CompareEntity(float* prevModelMatrix, int uniqueId, const float* modelMatrix)
	{
		Q_assert(begun);

		if(!enabled || uniqueId == 0)
		{
			memcpy(prevModelMatrix, modelMatrix, 16 * sizeof(float));
			return;
		}

		EntityData* prevEnt = NULL;
		for(uint32_t i = 0, end = prevEnts->count; i < end; i++)
		{
			EntityData& ent = (*prevEnts)[i];
			if(ent.uniqueId == uniqueId)
			{
				prevEnt = &ent;
				break;
			}
		}
		if(prevEnt != NULL)
		{
			memcpy(prevModelMatrix, prevEnt->modelMatrix, 16 * sizeof(float));
		}
		else
		{
			memcpy(prevModelMatrix, modelMatrix, 16 * sizeof(float));
		}

		EntityData* currEnt = NULL;
		for(uint32_t i = 0, end = currEnts->count; i < end; i++)
		{
			EntityData& ent = (*currEnts)[i];
			if(ent.uniqueId == uniqueId)
			{
				currEnt = &ent;
				break;
			}
		}
		if(currEnt != NULL)
		{
			memcpy(currEnt->modelMatrix, modelMatrix, 16 * sizeof(float));
		}
		else
		{
			EntityData ent;
			memcpy(ent.modelMatrix, modelMatrix, 16 * sizeof(float));
			ent.uniqueId = uniqueId;
			currEnts->Add(ent);
		}
	}

private:
	typedef RHI::StaticArray<EntityData, 16384> EntityArray;
	EntityArray entities[2];
	EntityArray* currEnts = &entities[0];
	EntityArray* prevEnts = &entities[1];
	bool enabled;
	bool begun;
};

static EntityTracker et;


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

	{
		GraphicsPipelineDesc desc("Skybox Motion Vectors");
		MakeFullScreenPipeline(desc, ShaderByteCode(g_skybox_motion_ps));
		desc.AddRenderTarget(0, TextureFormat::R16G16_Float);
		skyboxMotionPipeline = CreateGraphicsPipeline(desc);
	}
}

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;

	et.BeginFrame();

	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.motionVectorMBTexture, ResourceStates::RenderTargetBit);
	CmdTextureBarrier(crp.shadingPositionTexture, ResourceStates::RenderTargetBit);
	CmdEndBarrier();

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

	// clear the motion vectors with skybox values
	CmdBeginDebugLabel("Skybox Motion Vectors");
	{
		CmdBindRenderTargets(1, &crp.motionVectorTexture, NULL);
		CmdBindPipeline(skyboxMotionPipeline);
		CmdDraw(3, 0);

		if(crp_mblur->integer & MotionBlurModes::CameraBit)
		{
			CmdBindRenderTargets(1, &crp.motionVectorMBTexture, NULL);
			CmdBindPipeline(skyboxMotionPipeline);
			CmdDraw(3, 0);
		}
		else
		{
			CmdClearColorTarget(crp.motionVectorMBTexture, vec4_zero, NULL);
		}
	}
	CmdEndDebugLabel();

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

	const HTexture renderTargets[] =
	{
		crp.normalTexture,
		crp.motionVectorTexture,
		crp.motionVectorMBTexture,
		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;
	batchEntityId = -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;
		const bool idChanged = drawSurf->uniqueEntityId != batchEntityId;
		if(shaderChanged || entityChanged || idChanged)
		{
			oldShader = shader;
			oldEntityNum = entityNum;
			EndBatch();
			BeginBatch(shader);
			tess.greyscale = drawSurf->greyscale;
			batchEntityId = drawSurf->uniqueEntityId;
			batchMotionScale = drawSurf->motionBlurScale;
		}

		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;

	et.EndFrame();
}

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!
	if(rhiInfo.hasBarycentrics)
	{
		desc.vertexShader.Set(g_prepass_bary_vs);
		desc.pixelShader.Set(g_prepass_bary_ps);
	}
	else
	{
		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::R16G16_SNorm);
	desc.AddRenderTarget(0, TextureFormat::R16G16_Float);
	desc.AddRenderTarget(0, TextureFormat::R16G16_Float);
	desc.AddRenderTarget(0, TextureFormat::R32G32B32A32_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];
	float normalMatrix[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));

	R_InvMatrix(backEnd.modelMatrix, tempMatrix);
	R_TransposeMatrix(tempMatrix, normalMatrix);

	// select the right motion vector code path in the shader for view weapons
	// we should identify view weapons directly instead of using the depth hack flag
	int mblurMode = crp_mblur->integer;
	float mblurScale = batchMotionScale;
	if(batchDepthHack && mblurMode == MotionBlurModes::CameraOnly)
	{
		mblurScale = 0.0f;
		mblurMode = 0; // doesn't really matter but shows intent in the root constants
	}
	else if(batchDepthHack && mblurMode == MotionBlurModes::ObjectOnly)
	{
		mblurMode = MotionBlurModes::Full;
	}

	memcpy(rc.modelViewMatrix, backEnd.orient.modelMatrix, sizeof(rc.modelViewMatrix));
	memcpy(rc.modelMatrix, backEnd.modelMatrix, sizeof(rc.modelMatrix));
	et.CompareEntity(rc.prevModelMatrix, batchEntityId, backEnd.modelMatrix);
	rc.normalMatrix[0] = normalMatrix[ 0];
	rc.normalMatrix[1] = normalMatrix[ 1];
	rc.normalMatrix[2] = normalMatrix[ 2];
	rc.normalMatrix[3] = normalMatrix[ 4];
	rc.normalMatrix[4] = normalMatrix[ 5];
	rc.normalMatrix[5] = normalMatrix[ 6];
	rc.normalMatrix[6] = normalMatrix[ 8];
	rc.normalMatrix[7] = normalMatrix[ 9];
	rc.normalMatrix[8] = normalMatrix[10];
	rc.motionBlurScale = mblurScale;
	rc.textureIndex = texIdx;
	rc.samplerIndex = sampIdx;
	rc.alphaTest = alphaTest;
	rc.motionBlurMode = (uint32_t)mblurMode;
	CmdSetGraphicsRootConstants(0, sizeof(rc), &rc);

	db.DrawStage(vertexCount, indexCount);

	db.EndBaseBatch(vertexCount);

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