/*
===========================================================================
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_sunlight;
extern cvar_t* crp_volLight;
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 TranspResolve
{
	void Init();
	void Draw(const drawSceneViewCommand_t& cmd);

private:
	HPipeline noVolPipeline;
	HPipeline volPipeline;
};

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,
			SunlightVisibility,
			SunlightPenumbra,
			Sunlight,
			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 DrawBegin();
	void DrawPointLight(const dlight_t& light);
	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 SunlightEditor
{
	void Init();
	void ProcessWorld(world_t& world);
	void DrawOverlay();
	void DrawGUI();

private:
	HPipeline pipeline;
	bool windowActive = false;
	bool drawOverlay = false;
	const shader_t* skyShader = NULL;
};

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

//private:
	HPipeline visibilityPipeline;
	HPipeline blurPipeline;
	HTexture visibilityTexture;
	HTexture penumbraTexture;
};

struct VolumetricLight
{
	void Init();
	void ProcessWorld(world_t& world);
	void DrawBegin();
	void DrawPointLight(const dlight_t& light);
	void DrawSunlight();
	void DrawEnd();
	void DrawDebug();
	void DrawGUI();
	bool WantRTASUpdate(const trRefdef_t& scene);
	bool ShouldDraw();
	bool ShouldDrawDebug();
	bool LoadFogFile(const char* filePath);
	void SaveFogFile(const char* filePath);

	// GUI/user-friendly data layout
	struct Fog
	{
		vec3_t scatterColor;
		vec3_t emissiveColor;
		vec3_t boxMin;
		vec3_t boxMax;
		float extinction;
		float albedo; // scatter / extinction
		float emissive;
		float anisotropy;
		float noiseStrength;
		float noiseSpatialPeriod;
		float noiseTimePeriod;
		bool isGlobalFog;
		bool isHeightFog;
	};

#if defined(VL_CPU_PARTICLES)
	HPipeline particleDispatchPipeline;
	HPipeline particlePreProcessExtinctionPipeline;
	HPipeline particlePreProcessFrustumPipeline;
	HPipeline extinctionParticlePipeline;
	HPipeline frustumParticlePipeline;
	HBuffer particleBuffer; // should be double buffered...
	HBuffer particleHitBuffer; // for each tile, is there at least 1 particle?
	HBuffer particleTileBuffer; // array of tiles, each tile has a uint3 index
	HBuffer particleDispatchBuffer; // indirect dispatch buffer
	HBuffer particleDispatchClearBuffer; // indirect dispatch buffer with values (0, 1, 1)
	uint32_t particleCount;
#endif

	Fog fogs[64];
	uint32_t fogCount = 0;
	HPipeline extinctionFogPipeline;
	HPipeline frustumAmbientPipeline;
	HPipeline frustumAnisotropyPipeline;
	HPipeline frustumFogPipeline;
	HPipeline frustumPointLightScatterPipeline;
	HPipeline frustumRaymarchPipeline;
	HPipeline frustumSunlightVisPipeline;
	HPipeline frustumTemporalPipeline;
	HPipeline pointLightShadowPipeline;
	HPipeline sunlightScatterPipeline;
	HPipeline sunlightShadowPipeline;
	HPipeline ambientVizPipeline;
	HPipeline extinctionVizPipeline;
	HPipeline sunShadowVizPipeline;
	HTexture materialTextureA; // frustum, RGB = scatter, A = absorption
	HTexture materialTextureB; // frustum, RGB = emissive, A = anisotropy (g)
	HTexture materialTextureC; // frustum, R = anisotropy (g) weight sum
	HTexture sunlightVisTexture; // frustum, R = sunlight visibility
	HTexture prevSunlightVisTexture; // frustum, R = sunlight visibility
	HTexture scatterExtTexture; // frustum, RGB = in-scattering, A = extinction
	HTexture scatterTransTexture; // frustum, RGB = in-scattering, A = transmittance
	HTexture extinctionTextures[4]; // cube, R = extinction
	HTexture pointShadowTexture; // cube, R = transmittance
	HTexture sunShadowTextures[4]; // cube, R = transmittance
	HTexture ambientLightTextureA; // box, can be NULL, RGB = ambient.rgb, A = directional.r
	HTexture ambientLightTextureB; // box, can be NULL, RG = directional.gb, B = longitude, A = latitude
	uvec3_t frustumSize; // frustum volume pixel counts
	uvec3_t frustumTileScale; // by how much do we divide
	uvec3_t frustumTileSize; // frustum volume tile pixel counts
	uvec3_t extinctionSize; // extinction volume pixel counts
	uvec3_t extinctionTileScale; // by how much do we divide
	uvec3_t extinctionTileSize; // extinction volume tile pixel counts
	uint32_t shadowPixelCount; // @TODO: transform into uvec3_t as well
	uint32_t jitterCounter;
	uint32_t depthMip; // has to match the X/Y scale of frustumSize
	vec4_t extinctionVolumeScale; // how many world units per pixel
	float pointShadowVolumeScale; // how many world units per pixel
	vec4_t sunShadowVolumeScale; // how many world units per pixel
	uvec3_t sunShadowSize; // sunlight shadow volume pixel counts
	vec3_t ambientColor;
	float ambientIntensity;
	vec3_t debugCameraPosition;
	float debugBoxScale = 1.0f;
	float debugExtinctionScale = 50.0f;
	int debugExtinctionCascadeIndex = 0;
	int debugSunShadowCascadeIndex = 0;
	bool drawExtinctionDebug = false;
	bool drawSunShadowDebug = false;
	bool drawAmbientDebug = false;
	bool lockCameraPosition = false;
	bool firstFrame = true;
	bool windowActive = false;
	vec3_t mapBoxMin;
	vec3_t mapBoxMax;
	vec3_t lightGridCenter;
	float debugSphereScale = 0.5f;
};

#pragma pack(push, 1)
struct SunlightData
{
	vec3_t direction;
	vec3_t color;
	float intensity;
};
#pragma pack(pop)

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 DrawSceneView3D(const drawSceneViewCommand_t& cmd);
	void UploadSceneViewData();
	void BuildDepthPyramid();

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

	// general
	float frameSeed;
	HTexture readbackRenderTarget;
	HTexture depthTexture;
	HTexture depthMinMaxTexture;
	HTexture normalTexture;
	HTexture motionVectorTexture; // raw, for TAA/denoisers/etc
	HTexture motionVectorMBTexture; // mangled, for motion blur only
	HTexture sunlightTexture;
	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;
	HPipeline depthPyramidPipeline;

	// 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;
	Magnifier magnifier;
	DynamicLights dynamicLights;
	Sunlight sunlight;
	VolumetricLight volumetricLight;
	Raytracing raytracing;
	GBufferViz gbufferViz;
	SunlightEditor sunlightEditor;
	SunlightData sunlightData;
};

HPipeline CreateComputePipeline(const char* name, const ShaderByteCode& shader);
void MakeFullScreenPipeline(GraphicsPipelineDesc& desc, const ShaderByteCode& pixelShader);
void DirectionToAzimuthInclination(float* sc, const float* dir);
void AzimuthInclinationToDirection(float* dir, const float* sc);

extern CRP crp;