/*
===========================================================================
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 - GPU particle system


#include "crp_local.h"
#include "shaders/crp/scene_view.h.hlsli"
#include "compshaders/crp/particles_clear.h"
#include "compshaders/crp/particles_setup.h"
#include "compshaders/crp/particles_emit.h"
#include "compshaders/crp/particles_simulate.h"


 // @TODO:
#if 0
static const uint32_t EmitterParticleCount = 5;
static const uint32_t EmitterEmitCount = 2;
static const float EmitterMaxSeconds = 2.0f / 60.0f;
#else
static const float EmitterMaxSeconds = 1.0f;
static const float EmitterFPS = 60.0f;
static const uint32_t EmitterParticleCount = 1024;
static const uint32_t EmitterEmitCount = EmitterParticleCount / (uint32_t)ceilf(EmitterMaxSeconds * EmitterFPS);
#endif


#pragma pack(push, 4)

struct GPSClearRC
{
	uint32_t emitterBufferIndex;
	uint32_t deadBufferIndex;
	uint32_t emitterIndex;
	uint32_t firstParticle;
	uint32_t particleCount;
	float maxSeconds;
};

struct GPSSetupRC
{
	uint32_t emitterBufferIndex;
	uint32_t indirectBufferIndex;
	uint32_t emitterIndex;
	uint32_t emitCount;
};

struct GPSEmitRC
{
	uint32_t particleBufferIndex;
	uint32_t liveBufferIndex;
	uint32_t deadBufferIndex;
	uint32_t emitterBufferIndex;
	uint32_t emitterIndex;
};

struct GPSSimulateRC
{
	uint32_t particleBufferIndex;
	uint32_t liveSrcBufferIndex;
	uint32_t liveDstBufferIndex;
	uint32_t deadBufferIndex;
	uint32_t emitterBufferIndex;
	uint32_t emitterIndex;
};

#pragma pack(pop)


void ParticleSystem::Init()
{
	Q_assert(EmitterParticleCount <= MAX_PARTICLES);

	clearPipeline = CreateComputePipeline("GPS Clear", ShaderByteCode(g_particles_clear_cs));
	setupPipeline = CreateComputePipeline("GPS Setup", ShaderByteCode(g_particles_setup_cs));
	emitPipeline = CreateComputePipeline("GPS Emit", ShaderByteCode(g_particles_emit_cs));
	simulatePipeline = CreateComputePipeline("GPS Simulate", ShaderByteCode(g_particles_simulate_cs));

	{
		BufferDesc desc("particles", MAX_PARTICLES * sizeof(Particle), ResourceStates::UnorderedAccessBit);
		desc.shortLifeTime = true;
		desc.structureByteCount = sizeof(Particle);
		particleBuffer = CreateBuffer(desc);
	}

	{
		BufferDesc desc("", MAX_PARTICLES * 4, ResourceStates::UnorderedAccessBit);
		desc.shortLifeTime = true;
		desc.structureByteCount = 4;
		desc.name = "live particles #1";
		liveBuffers[0] = CreateBuffer(desc);
		desc.name = "live particles #2";
		liveBuffers[1] = CreateBuffer(desc);
		desc.name = "dead particles";
		deadBuffer = CreateBuffer(desc);
	}

	{
		BufferDesc desc("particle emitters", MAX_PARTICLE_EMITTERS * sizeof(ParticleEmitter), ResourceStates::UnorderedAccessBit);
		desc.shortLifeTime = true;
		desc.structureByteCount = sizeof(ParticleEmitter);
		emitterBuffer = CreateBuffer(desc);
	}

	{
		const uint32_t dispatchData[] =
		{ 
			0, 1, 1,
			0, 1, 1
		};

		BufferDesc desc("particle dispatch", sizeof(dispatchData), ResourceStates::UnorderedAccessBit);
		desc.shortLifeTime = true;
		indirectBuffer = CreateBuffer(desc);

		uint8_t* const mapped = BeginBufferUpload(indirectBuffer);
		memcpy(mapped, dispatchData, sizeof(dispatchData));
		EndBufferUpload(indirectBuffer);
	}

	needsClearing = true;
	liveBufferReadIndex = 0;
}

void ParticleSystem::Draw()
{
#if 1 // @TODO: shouldn't be necessary once cgame is adding the emitters
	if(tr.world == NULL)
	{
		return;
	}
#endif

#if 0
	static int counter = 0;
	counter++;
	if(counter == 4)
	{
		counter = 0;
		needsClearing = true;
	}
#endif

	SCOPED_RENDER_PASS("Particles", 1.0f, 1.0f, 1.0f);

	if(needsClearing)
	{
		SCOPED_DEBUG_LABEL("Clear", 1.0f, 1.0f, 1.0f);

		CmdBeginBarrier();
		CmdBufferBarrier(emitterBuffer, ResourceStates::UnorderedAccessBit);
		CmdBufferBarrier(deadBuffer, ResourceStates::UnorderedAccessBit);
		CmdEndBarrier();

		GPSClearRC rc = {};
		rc.emitterBufferIndex = GetBufferIndexUAV(emitterBuffer);
		rc.deadBufferIndex = GetBufferIndexUAV(deadBuffer);
		rc.emitterIndex = 0;
		rc.firstParticle = 0;
		rc.particleCount = EmitterParticleCount;
		rc.maxSeconds = EmitterMaxSeconds;

		CmdBindPipeline(clearPipeline);
		CmdSetComputeRootConstants(0, sizeof(rc), &rc);
		CmdDispatch((EmitterParticleCount + 63) / 64, 1, 1);

		needsClearing = false;
	}

	{
		SCOPED_DEBUG_LABEL("Setup", 1.0f, 1.0f, 1.0f);

		CmdBeginBarrier();
		CmdBufferBarrier(emitterBuffer, ResourceStates::UnorderedAccessBit);
		CmdBufferBarrier(indirectBuffer, ResourceStates::UnorderedAccessBit);
		CmdEndBarrier();

		GPSSetupRC rc = {};
		rc.emitterBufferIndex = GetBufferIndexUAV(emitterBuffer);
		rc.indirectBufferIndex = GetBufferIndexUAV(indirectBuffer);
		rc.emitterIndex = 0; // @TODO:
		rc.emitCount = EmitterEmitCount; // @TODO:

		CmdBindPipeline(setupPipeline);
		CmdSetComputeRootConstants(0, sizeof(rc), &rc);
		CmdDispatch(1, 1, 1);
	}

	{
		SCOPED_DEBUG_LABEL("Emit", 1.0f, 1.0f, 1.0f);

		CmdBeginBarrier();
		CmdBufferBarrier(indirectBuffer, ResourceStates::IndirectDispatchBit);
		CmdBufferBarrier(emitterBuffer, ResourceStates::UnorderedAccessBit);
		CmdBufferBarrier(deadBuffer, ResourceStates::UnorderedAccessBit);
		CmdBufferBarrier(liveBuffers[liveBufferReadIndex], ResourceStates::UnorderedAccessBit);
		CmdBufferBarrier(particleBuffer, ResourceStates::UnorderedAccessBit);
		CmdEndBarrier();

		GPSEmitRC rc = {};
		rc.emitterBufferIndex = GetBufferIndexUAV(emitterBuffer);
		rc.deadBufferIndex = GetBufferIndexUAV(deadBuffer);
		rc.liveBufferIndex = GetBufferIndexUAV(liveBuffers[liveBufferReadIndex]);
		rc.particleBufferIndex = GetBufferIndexUAV(particleBuffer);
		rc.emitterIndex = 0; // @TODO:

		CmdBindPipeline(emitPipeline);
		CmdSetComputeRootConstants(0, sizeof(rc), &rc);
		CmdDispatchIndirect(indirectBuffer, 0);
	}

	{
		SCOPED_DEBUG_LABEL("Simulate", 1.0f, 1.0f, 1.0f);

		CmdBeginBarrier();
		CmdBufferBarrier(indirectBuffer, ResourceStates::IndirectDispatchBit);
		CmdBufferBarrier(emitterBuffer, ResourceStates::UnorderedAccessBit);
		CmdBufferBarrier(deadBuffer, ResourceStates::UnorderedAccessBit);
		CmdBufferBarrier(liveBuffers[liveBufferReadIndex], ResourceStates::UnorderedAccessBit);
		CmdBufferBarrier(liveBuffers[liveBufferReadIndex ^ 1], ResourceStates::UnorderedAccessBit);
		CmdBufferBarrier(particleBuffer, ResourceStates::UnorderedAccessBit);
		CmdEndBarrier();

		GPSSimulateRC rc = {};
		rc.particleBufferIndex = GetBufferIndexUAV(particleBuffer);
		rc.liveSrcBufferIndex = GetBufferIndexUAV(liveBuffers[liveBufferReadIndex]);
		rc.liveDstBufferIndex = GetBufferIndexUAV(liveBuffers[liveBufferReadIndex ^ 1]);
		rc.deadBufferIndex = GetBufferIndexUAV(deadBuffer);
		rc.emitterBufferIndex = GetBufferIndexUAV(emitterBuffer);
		rc.emitterIndex = 0; // @TODO:

		CmdBindPipeline(simulatePipeline);
		CmdSetComputeRootConstants(0, sizeof(rc), &rc);
		CmdDispatchIndirect(indirectBuffer, 12);

		liveBufferReadIndex ^= 1;
	}
}