/*
===========================================================================
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 - opaque surfaces


#include "crp_local.h"
#include "compshaders/crp/opaque.h"
#include "compshaders/crp/wireframe_normals.h"
#include "compshaders/crp/add_light.h"


#pragma pack(push, 4)

struct OpaqueVertexRC : WorldVertexRC
{
};

struct OpaquePixelRC
{
	// general
	uint32_t textureIndex;
	uint32_t samplerIndex;
	uint32_t shaderIndexBufferIndex;
	uint32_t alphaTest;
	uint32_t lightTextureIndex;
	uint32_t lightmapPass;
	float greyscale;

	// shader trace
	uint32_t shaderTrace; // shader index: 14 - enable: 1
	uint16_t centerPixelX;
	uint16_t centerPixelY;
};

struct AddLightVertexRC : WorldVertexRC
{
};

#pragma pack(pop)


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

	{
		GraphicsPipelineDesc desc("Debug Normals");
		desc.shortLifeTime = true;
		desc.vertexShader.Set(g_wireframe_normals_vs);
		desc.pixelShader.Set(g_wireframe_normals_ps);
		desc.vertexLayout.AddAttribute(0, ShaderSemantic::Position, DataType::Float32, 3, 0);
		desc.vertexLayout.AddAttribute(1, ShaderSemantic::Color, DataType::UNorm8, 4, 0);
		desc.depthStencil.depthStencilFormat = TextureFormat::Depth32_Float;
		desc.depthStencil.depthComparison = ComparisonFunction::GreaterEqual;
		desc.depthStencil.enableDepthTest = true;
		desc.depthStencil.enableDepthWrites = false;
		desc.rasterizer.cullMode = CT_TWO_SIDED;
		desc.rasterizer.polygonOffset = false;
		desc.rasterizer.clampDepth = false;
		desc.rasterizer.wireFrame = true;
		desc.AddRenderTarget(0, crp.renderTargetFormat);
		wireframeNormalsPipeline = CreateGraphicsPipeline(desc);
	}
}

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

	srp.renderMode = RenderMode::World;

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

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

	CmdBeginBarrier();
	CmdTextureBarrier(crp.depthTexture, ResourceStates::DepthReadBit);
	CmdTextureBarrier(crp.lightTexture, ResourceStates::PixelShaderAccessBit);
	CmdTextureBarrier(crp.sunlightTexture, ResourceStates::PixelShaderAccessBit);
	CmdBufferBarrier(srp.traceRenderBuffer, ResourceStates::UnorderedAccessBit);
	CmdEndBarrier();

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

	SCOPED_RENDER_PASS("Opaque", 1.0f, 0.5f, 0.5f);

	CmdBindRenderTargets(1, &crp.renderTarget, &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);

		// sky shaders can have no stages and be valid (box drawn with no clouds)
		if(!shader->isSky)
		{
			if(shader->numPipelines == 0 ||
				shader->pipelines[0].pipeline <= 0 ||
				shader->pipelines[0].numStages <= 0)
			{
				continue;
			}
		}

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

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

		R_TessellateSurface(drawSurf->surface);
	}

	backEnd.refdef.floatTime = originalTime;

	EndSkyBatch();
	EndBatch();

	db.EndUpload();

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

void WorldOpaque::ProcessShader(shader_t& shader)
{
	Q_assert(shader.isOpaque || shader.isSky);

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

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

	for(int s = 0; s < shader.numStages; ++s)
	{
		const shaderStage_t& stage = *shader.stages[s];
		const unsigned int stateBits = stage.stateBits & (~GLS_POLYMODE_LINE);
		int a = 0;

		// @NOTE: we 0-init the struct so that padding bytes don't mess up comparisons in the PSO cache
		GraphicsPipelineDesc desc = {};
		desc.name = "opaque";
		desc.rootSignature = RHI_MAKE_NULL_HANDLE();
		desc.shortLifeTime = true; // the PSO cache is only valid for this map!
		desc.vertexShader.Set(g_opaque_vs);
		desc.pixelShader.Set(g_opaque_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 = shader.isSky ? ComparisonFunction::GreaterEqual : ComparisonFunction::Equal;
		desc.depthStencil.enableDepthTest = true;
		desc.depthStencil.enableDepthWrites = false;
		desc.rasterizer.cullMode = shader.cullType;
		desc.rasterizer.polygonOffset = shader.polygonOffset != 0;
		desc.rasterizer.clampDepth = clampDepth;
		desc.AddRenderTarget(stateBits & GLS_BLEND_BITS, crp.renderTargetFormat);

		pipeline_t& p = shader.pipelines[s];
		p.firstStage = s;
		p.numStages = 1;
		p.pipeline = psoCache.AddPipeline(desc, va("opaque %d %d", psoCache.entryCount, s + 1));
		desc.rasterizer.cullMode = GetMirrorredCullType(desc.rasterizer.cullMode);
		p.mirrorPipeline = psoCache.AddPipeline(desc, va("opaque %d %d mirrored", psoCache.entryCount, s + 1));
	}

	if(!shader.hasLightmapStage)
	{
		static int counter = 0;
		GraphicsPipelineDesc desc = {};
		desc.name = "Add Light";
		desc.rootSignature = RHI_MAKE_NULL_HANDLE();
		desc.shortLifeTime = true;
		desc.vertexShader.Set(g_add_light_vs);
		desc.pixelShader.Set(g_add_light_ps);
		desc.vertexLayout.AddAttribute(0, ShaderSemantic::Position, DataType::Float32, 3, 0);
		desc.depthStencil.depthStencilFormat = TextureFormat::Depth32_Float;
		desc.depthStencil.depthComparison = shader.isSky ? ComparisonFunction::GreaterEqual : ComparisonFunction::Equal;
		desc.depthStencil.enableDepthTest = true;
		desc.depthStencil.enableDepthWrites = false;
		desc.rasterizer.cullMode = shader.cullType;
		desc.rasterizer.polygonOffset = shader.polygonOffset != 0;
		desc.rasterizer.clampDepth = clampDepth;
		desc.AddRenderTarget(GLS_SRCBLEND_DST_COLOR | GLS_DSTBLEND_ONE, crp.renderTargetFormat);
		shader.addLightPipeline = psoCache.AddPipeline(desc, va("Add Light #%d", counter++ + 1));
	}

	shader.numPipelines = shader.numStages;
}

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

void WorldOpaque::BeginBatch(const shader_t* shader)
{
	tess.tessellator = Tessellator::Opaque;
	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 WorldOpaque::EndBatch()
{
	const int vertexCount = tess.numVertexes;
	const int indexCount = tess.numIndexes;
	if(vertexCount <= 0 ||
		indexCount <= 0 ||
		tess.shader->numStages == 0 ||
		tess.shader->numPipelines <= 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;
	}

	OpaqueVertexRC vertexRC = {};
	memcpy(vertexRC.modelViewMatrix, backEnd.orient.modelMatrix, sizeof(vertexRC.modelViewMatrix));
	CmdSetGraphicsRootConstants(0, sizeof(vertexRC), &vertexRC);

	for(int s = 0; s < shader->numStages; ++s)
	{
		const shaderStage_t* const stage = shader->stages[s];

		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->pipelines[s];
		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);
		const uint32_t enableShaderTrace = tr.traceWorldShader && s == 0 ? 1 : 0;
		const uint32_t bufferIndex = GetBufferIndexUAV(srp.traceRenderBuffer);
		Q_assert(sampIdx < ARRAY_LEN(crp.samplers));

		OpaquePixelRC pixelRC = {};
		pixelRC.textureIndex = texIdx;
		pixelRC.samplerIndex = sampIdx;
		pixelRC.shaderIndexBufferIndex = bufferIndex;
		pixelRC.alphaTest = alphaTest;
		pixelRC.lightTextureIndex = GetTextureIndexSRV(crp.lightTexture);
		pixelRC.lightmapPass = stage->type == ST_LIGHTMAP ? 1 : 0;
		pixelRC.greyscale = tess.greyscale;
		pixelRC.shaderTrace = ((uint32_t)shader->index << 1) | enableShaderTrace;
		pixelRC.centerPixelX = glConfig.vidWidth / 2;
		pixelRC.centerPixelY = glConfig.vidHeight / 2;
		CmdSetGraphicsRootConstants(sizeof(vertexRC), sizeof(pixelRC), &pixelRC);

		db.DrawStage(vertexCount, indexCount);
	}

	if(!shader->hasLightmapStage)
	{
		AddLightVertexRC rc = {};
		memcpy(rc.modelViewMatrix, backEnd.orient.modelMatrix, sizeof(rc.modelViewMatrix));
		CmdSetGraphicsRootConstants(0, sizeof(rc), &rc);
		CmdBindPipeline(psoCache.entries[shader->addLightPipeline].handle);
		db.DrawPositionOnly(vertexCount, indexCount);
	}

	db.EndBaseBatch(vertexCount);

	if(crp_drawNormals->integer)
	{
		CmdBindPipeline(wireframeNormalsPipeline);
		db.UploadAndDrawDebugNormals();
	}

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

void WorldOpaque::EndSkyBatch()
{
	// this only exists as a separate function from EndBatch so that
	// we don't have to deal with recursion (through the call to RB_DrawSky)

	if(tess.shader == NULL ||
		!tess.shader->isSky ||
		tess.numVertexes <= 0 ||
		tess.numIndexes <= 0)
	{
		return;
	}

	SCOPED_RENDER_PASS("Sky", 0.0, 0.5f, 1.0f);

	const viewParms_t& vp = backEnd.viewParms;
	CmdSetViewport(vp.viewportX, vp.viewportY, vp.viewportWidth, vp.viewportHeight, 0.0f, 0.0f);
	RB_DrawSky();
	CmdSetViewport(vp.viewportX, vp.viewportY, vp.viewportWidth, vp.viewportHeight, 0.0f, 1.0f);
	tess.numVertexes = 0;
	tess.numIndexes = 0;
}

void WorldOpaque::DrawSkyBox()
{
	// force creation of a PSO for the temp shader
	ProcessShader((shader_t&)*tess.shader);

	tess.deformsPreApplied = qtrue;
	EndBatch();
}

void WorldOpaque::DrawClouds()
{
	EndBatch();
}