/*
===========================================================================
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 - private declarations


#pragma once


#include "srp_local.h"


extern cvar_t* crp_dof;
extern cvar_t* crp_dof_overlay;
extern cvar_t* crp_dof_blades;
extern cvar_t* crp_dof_angle;
extern cvar_t* crp_gatherDof_focusNearDist;
extern cvar_t* crp_gatherDof_focusNearRange;
extern cvar_t* crp_gatherDof_focusFarDist;
extern cvar_t* crp_gatherDof_focusFarRange;
extern cvar_t* crp_gatherDof_brightness;
extern cvar_t* crp_accumDof_focusDist;
extern cvar_t* crp_accumDof_radius;
extern cvar_t* crp_accumDof_samples;
extern cvar_t* crp_accumDof_preview;
extern cvar_t* crp_mblur;
extern cvar_t* crp_mblur_exposure;
extern cvar_t* crp_drawNormals;
extern cvar_t* crp_updateRTAS;
extern cvar_t* crp_debug0;
extern cvar_t* crp_debug1;
extern cvar_t* crp_debug2;
extern cvar_t* crp_debug3;

struct DOFMethod
{
	enum Id
	{
		None,
		Gather,
		Accumulation,
		Count
	};
};

struct MotionBlurModes
{
	enum Flags
	{
		CameraBit = 1 << 0,
		ObjectBit = 1 << 1
	};
};

struct Tessellator
{
	enum Id
	{
		None,
		Prepass,
		Opaque,
		Transp,
		Count
	};
};


using namespace RHI;


struct WorldVertexRC
{
	float modelViewMatrix[16];
};

struct PSOCache
{
	struct Entry
	{
		GraphicsPipelineDesc desc;
		HPipeline handle;
	};

	void Init(Entry* entries, uint32_t maxEntryCount);
	int AddPipeline(const GraphicsPipelineDesc& desc, const char* name);
	
	Entry* entries = NULL;
	uint32_t maxEntryCount = 0;
	uint32_t entryCount = 1; // we treat index 0 as invalid
};

struct Prepass
{
	void Init();
	void Draw(const drawSceneViewCommand_t& cmd);
	void ProcessShader(shader_t& shader);
	void TessellationOverflow();

private:
	void BeginBatch(const shader_t* shader);
	void EndBatch();

	PSOCache::Entry psoCacheEntries[128];
	PSOCache psoCache;

	bool batchOldDepthHack;
	bool batchDepthHack;
	int batchEntityId;
	float batchMotionScale;

	HPipeline skyboxMotionPipeline;
};

struct WorldOpaque
{
	void Init();
	void Draw(const drawSceneViewCommand_t& cmd);
	void ProcessShader(shader_t& shader);
	void TessellationOverflow();
	void DrawSkyBox();
	void DrawClouds();

private:
	void BeginBatch(const shader_t* shader);
	void EndBatch();
	void EndSkyBatch();

	PSOCache::Entry psoCacheEntries[128];
	PSOCache psoCache;

	bool batchOldDepthHack;
	bool batchDepthHack;

	HPipeline wireframeNormalsPipeline;
};

struct WorldTransp
{
	void Init();
	void Draw(const drawSceneViewCommand_t& cmd);
	void ProcessShader(shader_t& shader);
	void TessellationOverflow();

private:
	void BeginBatch(const shader_t* shader);
	void EndBatch();

	PSOCache::Entry psoCacheEntries[32];
	PSOCache psoCache;

	bool batchOldDepthHack;
	bool batchDepthHack;
};

struct Fog
{
	void Init();
	void Draw();

private:
	HBuffer boxIndexBuffer;
	HBuffer boxVertexBuffer;
	HPipeline fogInsidePipeline;
	HPipeline fogOutsidePipeline;
};

struct TranspResolve
{
	void Init();
	void Draw(const drawSceneViewCommand_t& cmd);

private:
	HPipeline pipeline;
};

struct ToneMap
{
	void Init();
	void DrawToneMap();
	void DrawInverseToneMap();

private:
	HPipeline pipeline;
	HPipeline inversePipeline;
};

struct AccumDepthOfField
{
	void Init();
	void Begin(const drawSceneViewCommand_t& cmd);
	uint32_t GetSampleCount();
	void FixCommand(drawSceneViewCommand_t& newCmd, const drawSceneViewCommand_t& cmd, uint32_t x, uint32_t y);
	void Accumulate();
	void Normalize();
	void DrawDebug();

private:
	HPipeline accumPipeline;
	HPipeline normPipeline;
	HPipeline debugPipeline;
	HTexture accumTexture;
	float maxNearCocCS;
	float maxFarCocCS;
	float modelViewMatrix[16];
	float projMatrix[16];
};

struct GatherDepthOfField
{
	void Init();
	void Draw();

private:
	void DrawDebug();
	void DrawSplit();
	void DrawNearCocTileGen();
	void DrawNearCocTileMax();
	void DrawBlur();
	void DrawFill();
	void DrawCombine();

	HPipeline debugPipeline;
	HPipeline splitPipeline;
	HPipeline nearCocTileGenPipeline;
	HPipeline nearCocTileMaxPipeline;
	HPipeline blurPipeline;
	HPipeline fillPipeline;
	HPipeline combinePipeline;
	HTexture nearColorTexture;
	HTexture farColorTexture;
	HTexture nearBlurTexture;
	HTexture farBlurTexture;
	HTexture nearCocTexture;
	HTexture nearCocTexture2;
	HTexture nearCocTileTexture;
	HTexture nearCocTileTexture2;
	HTexture farCocTexture;
	uint32_t tileTextureWidth;
	uint32_t tileTextureHeight;
};

struct MotionBlur
{
	void Init();
	void Draw();

private:
	void DrawPack();
	void DrawTileGen();
	void DrawTileMax();
	void DrawBlur();

	HPipeline packPipeline;
	HPipeline tileGenPipeline;
	HPipeline tileMaxPipeline;
	HPipeline blurPipeline;
	HTexture tileTexture;
	HTexture tileTexture2;
	HTexture packedTexture;
	uint32_t tileTextureWidth;
	uint32_t tileTextureHeight;
};

struct Magnifier
{
	void Init();
	void Draw();
	void DrawGUI();

private:
	HPipeline pipeline;
	bool magnifierActive = false;
};

struct GBufferViz
{
	void Init();
	void DrawGUI();

private:
	struct GBufferTexture
	{
		enum Id
		{
			Depth,
			Normal,
			Light,
			ShadingPositionDelta,
			MotionVectorRaw,
			MotionVectorMB,
			Count
		};
	};

	HPipeline linearizeDepthPipeline;
	HPipeline decodeNormalsPipeline;
	HPipeline decodeShadingPositionPipeline;
	HPipeline motionVectorPipeline;
	bool windowActive = false;
	int textureIndex = 0;
	bool coloredPositionDelta = false;
	bool fullResolution = false;
};

struct DynamicLights
{
	void Init();
	void Draw();
	bool WantRTASUpdate(const trRefdef_t& scene);

private:
	HPipeline pipeline;
	HPipeline denoisingPipeline;
};

struct Raytracing
{
	void Init();
	void ProcessWorld(world_t& world);
	void BeginFrame(bool wantUpdate);
	uint32_t GetTLASBufferIndex();
	uint32_t GetInstanceBufferIndex();
	bool CanRaytrace();

private:
	void TagMapSurfacesRecursively(mnode_t* node);

	struct BLASBucket
	{
		enum Constants
		{
			Count = CT_COUNT
		};
	};

	struct BLASBuildBuffers
	{
		HBuffer vertexBuffer;
		HBuffer indexBuffer;
		uint32_t vertexBufferByteCount;
		uint32_t indexBufferByteCount;
	};

	struct BLASBuffers
	{
		HBuffer blasBuffer;
		HBuffer vertexBuffer;
		HBuffer indexBuffer;
		HBuffer meshBuffer;
		uint32_t vertexBufferByteCount;
		uint32_t indexBufferByteCount;
		uint32_t meshBufferByteCount;
	};

	struct Surface
	{
		const surfaceType_t* surface;
		const shader_t* shader;
		int entityNum;
	};

	struct ISurfaceList
	{
		virtual uint32_t GetSurfaceCount() = 0;
		virtual bool GetSurface(Surface& surface, uint32_t index) = 0; // true when skipped
	};

	struct WorldSurfaceList : ISurfaceList
	{
		uint32_t GetSurfaceCount() override;
		bool GetSurface(Surface& surface, uint32_t index) override;
	};

	struct DynamicSurfaceList : ISurfaceList
	{
		uint32_t GetSurfaceCount() override;
		bool GetSurface(Surface& surface, uint32_t index) override;
	};

	void EnsureBuffersAreLargeEnough(BLASBuildBuffers& buffers, uint32_t maxVertexCount, uint32_t maxIndexCount);
	void EnsureBuffersAreLargeEnough(BLASBuffers& buffers, uint32_t maxVertexCount, uint32_t maxIndexCount, uint32_t maxMeshCount);
	void BuildBLASes(BLASBuffers* blasBuffers, struct BLASBuilder* blasBuilders, ISurfaceList* surfaceList);

	struct FrameData
	{
		BLASBuildBuffers blasBuildBuffers[BLASBucket::Count] = {};
		BLASBuffers dynamicBLASBuffers[BLASBucket::Count] = {};
		HBuffer tlasBuffer = RHI_MAKE_NULL_HANDLE();
		HBuffer tlasInstanceBuffer = RHI_MAKE_NULL_HANDLE();
	};

	enum Constants { RTFrameCount = RHI::FrameCount + 1 };

	uint32_t GetRTFrameIndex() { return tr.frameCount % RTFrameCount; }

	FrameData frameData[RTFrameCount];
	BLASBuffers staticBLASBuffers[BLASBucket::Count] = {};
	StaticUnorderedArray<TLASInstanceDesc, 2 * BLASBucket::Count> tlasInstanceDescs;
	uint32_t staticTLASInstanceCount = 0;
};

struct BaseBufferId
{
	enum Id
	{
		Position,
		Normal,
		Count
	};
};

struct StageBufferId
{
	enum Id
	{
		TexCoords,
		Color,
		Count
	};
};

struct GeoBuffers
{
	void Create(const char* name, uint32_t vertexCount, uint32_t indexCount);
	void Rewind();
	void BeginUpload();
	void EndUpload();
	void UploadBase();
	void UploadStage(uint32_t svarsIndex);
	void EndBaseBatch(uint32_t vertexCount);
	bool CanAdd(uint32_t vertexCount, uint32_t indexCount, uint32_t stageCount);
	void DrawStage(uint32_t vertexCount, uint32_t indexCount);
	void DrawPositionOnly(uint32_t vertexCount, uint32_t indexCount);
	void UploadAndDrawDebugNormals();

	GeometryBuffer baseVertexBuffers[BaseBufferId::Count];
	GeometryBuffer stageVertexBuffers[StageBufferId::Count];
	IndexBuffer indexBuffer;
	HBuffer vertexBuffers[BaseBufferId::Count + StageBufferId::Count];
	uint32_t vertexBufferStrides[BaseBufferId::Count + StageBufferId::Count];
};

struct FreezeFrame
{
	enum Id
	{
		Inactive,
		Pending,
		Active,
		PendingBeforeMB,
		ActiveBeforeMB
	};
};

struct CRP : IRenderPipeline
{
	void Init() override;
	void LoadResources() override;
	void ShutDown(bool fullShutDown) override;

	void ProcessWorld(world_t& world) override;
	void ProcessModel(model_t& model) override;
	void ProcessShader(shader_t& shader) override;

	void CreateTexture(image_t* image, int mipCount, int width, int height) override;
	void UpoadTextureAndGenerateMipMaps(image_t* image, const byte* data) override;
	void BeginTextureUpload(MappedTexture& mappedTexture, image_t* image) override;
	void EndTextureUpload() override;
	
	void ExecuteRenderCommands(const byte* data, bool readbackRequested) override;
	void TessellationOverflow() override;
	void DrawSkyBox() override { opaque.DrawSkyBox(); }
	void DrawClouds() override { opaque.DrawClouds(); }
	void ReadPixels(int w, int h, int alignment, colorSpace_t colorSpace, void* out) override;

	uint32_t GetSamplerDescriptorIndexFromBaseIndex(uint32_t baseIndex) override;

	void BeginFrame();
	void EndFrame();

	void Blit(HTexture destination, HTexture source, const char* passName, bool hdr, const vec2_t tcScale, const vec2_t tcBias);
	void BlitRenderTarget(HTexture destination, const char* passName);
	void DrawSceneView(const drawSceneViewCommand_t& cmd);
	void UploadSceneViewData();

	HTexture GetReadRenderTarget();
	HTexture GetWriteRenderTarget();
	void SwapRenderTargets();

	// general
	float frameSeed;
	HTexture readbackRenderTarget;
	HTexture depthTexture;
	HTexture normalTexture;
	HTexture motionVectorTexture; // raw, for TAA/denoisers/etc
	HTexture motionVectorMBTexture; // mangled, for motion blur only
	HTexture noisyLightTexture;
	HTexture lightTexture;
	HTexture shadingPositionTexture;
	HTexture renderTarget;
	TextureFormat::Id renderTargetFormat;
	HTexture renderTargets[2];
	uint32_t renderTargetIndex; // the one to write to
	HSampler samplers[BASE_SAMPLER_COUNT]; // all base samplers
	uint32_t samplerIndices[BASE_SAMPLER_COUNT]; // descriptor heap indices
	HTexture blueNoise2D;
	FreezeFrame::Id freezeFrame;
	HTexture frozenTexture;

	// blit
	HPipeline blitPipelineLDR;
	HPipeline blitPipelineHDR;

	// world geometry
	GeoBuffers dynBuffers[FrameCount]; // for rendering world surfaces

	// scene view data
	HBuffer sceneViewUploadBuffers[FrameCount];
	HBuffer sceneViewBuffer; // this is the buffer that lives at ResourceDescriptorHeap[0]
	uint32_t sceneViewIndex;

	// for rendering transparent world surfaces
	HTexture oitIndexTexture;
	HBuffer oitFragmentBuffer;
	HBuffer oitCounterBuffer;
	HBuffer oitCounterStagingBuffer;

	UI ui;
	MipMapGenerator mipMapGen;
	ImGUI imgui;
	Nuklear nuklear;
	Prepass prepass;
	WorldOpaque opaque;
	WorldTransp transp;
	TranspResolve transpResolve;
	ToneMap toneMap;
	GatherDepthOfField gatherDof;
	AccumDepthOfField accumDof;
	MotionBlur motionBlur;
	Fog fog;
	Magnifier magnifier;
	DynamicLights dynamicLights;
	Raytracing raytracing;
	GBufferViz gbufferViz;
};

HPipeline CreateComputePipeline(const char* name, const ShaderByteCode& shader);
void MakeFullScreenPipeline(GraphicsPipelineDesc& desc, const ShaderByteCode& pixelShader);

extern CRP crp;
