/*
===========================================================================
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 - OIT geometry pass


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


#pragma pack(push, 4)

struct TranspDrawVertexRC : WorldVertexRC
{
};

struct TranspDrawPixelRC
{
	uint32_t textureIndex;
	uint32_t samplerIndex;
	uint32_t alphaTest;
	uint32_t counterBuffer;
	uint32_t indexTexture;
	uint32_t fragmentBuffer;
	float greyscale;
	uint32_t stateBits;
	uint32_t shaderTrace;
	uint16_t hFadeDistance;
	uint16_t hFadeOffset;
	uint32_t depthFadeScaleBiasPO; // polygon offset: 1 - enable: 1 - color bias: 4 - color scale: 4
};

#pragma pack(pop)


static uint32_t GetFixedStageBits(uint32_t stateBits, uint32_t stageIndex)
{
	// makes sure we're not overwriting anything useful
	assert((stateBits & GLS_STAGEINDEX_BITS) == 0);

	// transform "no blend" into a "replace" blend mode
	if((stateBits & (GLS_SRCBLEND_BITS | GLS_DSTBLEND_BITS)) == 0)
	{
		stateBits |= GLS_SRCBLEND_ONE | GLS_DSTBLEND_ZERO;
	}

	stateBits |= stageIndex << GLS_STAGEINDEX_SHIFT;

	return stateBits;
}


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

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

	srp.renderMode = RenderMode::World;

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

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

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

	CmdBeginBarrier();
	CmdTextureBarrier(crp.depthTexture, ResourceStates::DepthWriteBit);
	CmdTextureBarrier(crp.oitIndexTexture, ResourceStates::UnorderedAccessBit);
	CmdBufferBarrier(crp.oitFragmentBuffer, ResourceStates::UnorderedAccessBit);
	CmdBufferBarrier(crp.oitCounterBuffer, ResourceStates::CopyDestinationBit);
	CmdEndBarrier();

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

	CmdBindRenderTargets(0, NULL, &crp.depthTexture);
	CmdBindVertexBuffers(ARRAY_LEN(db.vertexBuffers), db.vertexBuffers, db.vertexBufferStrides, NULL);
	CmdBindIndexBuffer(db.indexBuffer.buffer, IndexType::UInt32, 0);

	// reset the fragment counter
	CmdCopyBuffer(crp.oitCounterBuffer, crp.oitCounterStagingBuffer);

	// clear the index texture
	const uint32_t zeroes[4] = {};
	CmdClearTextureUAV(crp.oitIndexTexture, 0, zeroes);

	CmdBeginBarrier();
	CmdBufferBarrier(crp.oitCounterBuffer, ResourceStates::UnorderedAccessBit);
	CmdTextureBarrier(crp.oitIndexTexture, ResourceStates::UnorderedAccessBit);
	CmdEndBarrier();

	const drawSurf_t* drawSurfs = cmd.drawSurfs;
	const int opaqueSurfCount = cmd.numDrawSurfs - cmd.numTranspSurfs;
	const int transpSurfCount = 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;

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

		if(shader->isFog)
		{
			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 WorldTransp::ProcessShader(shader_t& shader)
{
	Q_assert(!shader.isOpaque);

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

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

	for(int s = 0; s < shader.numStages; ++s)
	{
		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 = "transp";
		desc.rootSignature = RHI_MAKE_NULL_HANDLE();
		desc.shortLifeTime = true; // the PSO cache is only valid for this map!
		desc.vertexShader.Set(g_transp_draw_vs);
		desc.pixelShader.Set(g_transp_draw_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 = false;
		desc.rasterizer.cullMode = shader.cullType;
		desc.rasterizer.polygonOffset = shader.polygonOffset != 0;
		desc.rasterizer.clampDepth = clampDepth;

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

	shader.numTranspPipelines = shader.numStages;
}

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

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

	TranspDrawVertexRC 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->transpPipelines[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 enableDepthFade = shader->dfType != DFT_NONE ? 1 : 0;
		const uint32_t polygonOffset = shader->polygonOffset ? 1 : 0;
		Q_assert(sampIdx < ARRAY_LEN(crp.samplers));

		TranspDrawPixelRC pixelRC = {};
		pixelRC.alphaTest = alphaTest;
		pixelRC.counterBuffer = GetBufferIndexUAV(crp.oitCounterBuffer);
		pixelRC.fragmentBuffer = GetBufferIndexUAV(crp.oitFragmentBuffer);
		pixelRC.greyscale = tess.greyscale;
		pixelRC.indexTexture = GetTextureIndexUAV(crp.oitIndexTexture, 0);
		pixelRC.samplerIndex = sampIdx;
		pixelRC.stateBits = GetFixedStageBits(stage->stateBits, s);
		pixelRC.textureIndex = texIdx;
		pixelRC.shaderTrace = ((uint32_t)shader->index << 1) | enableShaderTrace;
		pixelRC.hFadeDistance = f32tof16(shader->dfInvDist);
		pixelRC.hFadeOffset = f32tof16(shader->dfBias);
		pixelRC.depthFadeScaleBiasPO = (polygonOffset << 9) | (enableDepthFade << 8) | (uint32_t)r_depthFadeScaleAndBias[shader->dfType];
		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;
}