/*
===========================================================================
Copyright (C) 2023 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"
namespace opaque
{
#include "compshaders/crp/opaque_vs.h"
#include "compshaders/crp/opaque_ps.h"
}


#pragma pack(push, 4)

struct OpaqueVertexRC : WorldVertexRC
{
};

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

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

#pragma pack(pop)


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

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;

	if(backEnd.viewParms.isPortal)
	{
		float plane[4];
		plane[0] = backEnd.viewParms.portalPlane.normal[0];
		plane[1] = backEnd.viewParms.portalPlane.normal[1];
		plane[2] = backEnd.viewParms.portalPlane.normal[2];
		plane[3] = backEnd.viewParms.portalPlane.dist;

		float plane2[4];
		plane2[0] = DotProduct(backEnd.viewParms.orient.axis[0], plane);
		plane2[1] = DotProduct(backEnd.viewParms.orient.axis[1], plane);
		plane2[2] = DotProduct(backEnd.viewParms.orient.axis[2], plane);
		plane2[3] = DotProduct(plane, backEnd.viewParms.orient.origin) - plane[3];

		float* o = plane;
		const float* m = s_flipMatrix;
		const float* v = plane2;
		o[0] = m[0] * v[0] + m[4] * v[1] + m[8] * v[2] + m[12] * v[3];
		o[1] = m[1] * v[0] + m[5] * v[1] + m[9] * v[2] + m[13] * v[3];
		o[2] = m[2] * v[0] + m[6] * v[1] + m[10] * v[2] + m[14] * v[3];
		o[3] = m[3] * v[0] + m[7] * v[1] + m[11] * v[2] + m[15] * v[3];

		memcpy(clipPlane, plane, sizeof(clipPlane));
	}
	else
	{
		memset(clipPlane, 0, sizeof(clipPlane));
	}

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

	TextureBarrier tb(crp.depthTexture, ResourceStates::DepthWriteBit);
	BufferBarrier bb(srp.traceRenderBuffer, ResourceStates::UnorderedAccessBit);
	CmdBarrier(1, &tb, 1, &bb);

	CmdClearDepthStencilTarget(crp.depthTexture, true, 0.0f);

	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 are not using any CTOR because we deliberately want to 0-init the struct
		// this is necessary for padding bytes not to 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 = opaque::g_vs;
		desc.pixelShader = opaque::g_ps;
		desc.vertexLayout.AddAttribute(a++, ShaderSemantic::Position, DataType::Float32, 3, 0);
		desc.vertexLayout.AddAttribute(a++, ShaderSemantic::Normal, DataType::Float32, 2, 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 =
			(stateBits & GLS_DEPTHFUNC_EQUAL) != 0 ?
			ComparisonFunction::Equal :
			ComparisonFunction::GreaterEqual;
		desc.depthStencil.enableDepthTest = (stateBits & GLS_DEPTHTEST_DISABLE) == 0;
		desc.depthStencil.enableDepthWrites = (stateBits & GLS_DEPTHMASK_TRUE) != 0;
		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));
	}

	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));
	memcpy(vertexRC.projectionMatrix, backEnd.viewParms.projectionMatrix, sizeof(vertexRC.projectionMatrix));
	memcpy(vertexRC.clipPlane, clipPlane, sizeof(vertexRC.clipPlane));
	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.greyscale = tess.greyscale;
		pixelRC.shaderTrace = ((uint32_t)shader->index << 3) | (RHI::GetFrameIndex() << 1) | enableShaderTrace;
		pixelRC.centerPixelX = glConfig.vidWidth / 2;
		pixelRC.centerPixelY = glConfig.vidHeight / 2;
		CmdSetGraphicsRootConstants(sizeof(vertexRC), sizeof(pixelRC), &pixelRC);

		db.DrawStage(vertexCount, indexCount);
	}

	db.EndBaseBatch(vertexCount);

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();
}