/*
===========================================================================
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/>.
===========================================================================
*/
// Shared Rendering Pipeline - private declarations


#pragma once


#include "tr_local.h"
#include "rhi_local.h"


using namespace RHI;


struct BufferBase
{
	bool CanAdd(uint32_t count_)
	{
		return batchFirst + batchCount + count_ <= totalCount;
	}

	void EndBatch()
	{
		batchFirst += batchCount;
		batchCount = 0;
	}

	void EndBatch(uint32_t size)
	{
		batchFirst += size;
		batchCount = 0;
	}

	void Rewind()
	{
		batchFirst = 0;
		batchCount = 0;
	}

	uint32_t totalCount = 0;
	uint32_t batchFirst = 0;
	uint32_t batchCount = 0;
};

struct IndexBuffer : BufferBase
{
	void Create(const char* name, MemoryUsage::Id memoryUsage, uint32_t indexCount)
	{
		totalCount = indexCount;

		BufferDesc desc = {};
		desc.committedResource = true;
		desc.initialState = ResourceStates::IndexBufferBit;
		desc.memoryUsage = memoryUsage;
		desc.name = va("%s index", name);
		desc.byteCount = indexCount * sizeof(uint32_t);
		buffer = CreateBuffer(desc);
	}

	void BeginUpload()
	{
		mapped = (uint32_t*)BeginBufferUpload(buffer);
	}

	void EndUpload()
	{
		EndBufferUpload(buffer);
		mapped = NULL;
	}

	void Upload()
	{
		Q_assert(mapped != NULL);

		uint32_t* const idx = mapped + batchFirst + batchCount;
		memcpy(idx, &tess.indexes[0], tess.numIndexes * sizeof(uint32_t));
	}

	uint32_t* GetCurrentAddress()
	{
		return mapped + batchFirst + batchCount;
	}

	HBuffer buffer = RHI_MAKE_NULL_HANDLE();
	uint32_t* mapped = NULL;
};

struct GeometryBuffer : BufferBase
{
	void Init(uint32_t count_, uint32_t stride_)
	{
		buffer = RHI_MAKE_NULL_HANDLE();
		byteCount = count_ * stride_;
		stride = stride_;
		totalCount = count_;
		batchFirst = 0;
		batchCount = 0;
	}

	void CreateVertexBuffer(const char* name, MemoryUsage::Id memoryUsage, uint32_t count, uint32_t stride_)
	{
		BufferDesc desc = {};
		desc.committedResource = true;
		desc.initialState = ResourceStates::VertexBufferBit;
		desc.memoryUsage = memoryUsage;
		desc.name = name;
		desc.byteCount = count * stride_;
		buffer = CreateBuffer(desc);
		byteCount = count * stride_;
		stride = stride_;
		totalCount = count;
		batchFirst = 0;
		batchCount = 0;
	}

	void BeginUpload()
	{
		Q_assert(mapped == NULL);
		mapped = BeginBufferUpload(buffer);
	}

	void EndUpload()
	{
		Q_assert(mapped != NULL);
		EndBufferUpload(buffer);
		mapped = NULL;
	}

	HBuffer buffer = RHI_MAKE_NULL_HANDLE();
	uint32_t byteCount = 0;
	uint32_t stride = 0;
	uint8_t* mapped = NULL;
};

struct RenderMode
{
	enum Id
	{
		None,
		UI,
		World,
		ImGui,
		Nuklear,
		Count
	};
};

struct RenderPassQueries
{
	char name[64];
	uint32_t gpuDurationUS;
	uint32_t cpuDurationUS;
	int64_t cpuStartUS;
	uint32_t queryIndex;
};

enum
{
	MaxRenderPasses = 64, // cg_draw3dIcons forces tons of 2D/3D transitions...
	MaxStatsFrameCount = 64
};

struct RenderPassStats
{
	void EndFrame(uint32_t cpu, uint32_t gpu);

	uint32_t samplesCPU[MaxStatsFrameCount];
	uint32_t samplesGPU[MaxStatsFrameCount];
	stats_t statsCPU;
	stats_t statsGPU;
	uint32_t count;
	uint32_t index;
};

struct RenderPassFrame
{
	RenderPassQueries passes[MaxRenderPasses];
	uint32_t count;
};

struct FrameStats
{
	enum { MaxFrames = 1024 };

	void EndFrame();

	float temp[MaxFrames];
	float p2pMS[MaxFrames];
	stats_t p2pStats;
	int frameCount;
	int frameIndex;
	int skippedFrames;
};

struct MipMapGenerator
{
	void Init(bool ddhi, const ShaderByteCode& g2l, const ShaderByteCode& down, const ShaderByteCode& l2g);
	void GenerateMipMaps(HTexture texture);

private:
	struct Stage
	{
		enum Id
		{
			Start,      // gamma to linear
			DownSample, // down sample on 1 axis
			End,        // linear to gamma
			Count
		};

		HRootSignature rootSignature;
		HDescriptorTable descriptorTable;
		HPipeline pipeline;
	};

	struct MipSlice
	{
		enum Id
		{
			Float16_0,
			Float16_1,
			Count
		};
	};

	HTexture textures[MipSlice::Count];
	Stage stages[3];
	bool ddhi = false;
};

struct UI
{
	void Init(bool ddhi_, const ShaderByteCode& vs, const ShaderByteCode& ps, TextureFormat::Id rtFormat,
		HDescriptorTable descTable, RootSignatureDesc* rootSigDesc);
	void BeginFrame();
	void Begin(HTexture renderTarget);
	void End();
	void CmdSetColor(const uiSetColorCommand_t& cmd);
	void CmdDrawQuad(const uiDrawQuadCommand_t& cmd);
	void CmdDrawTriangle(const uiDrawTriangleCommand_t& cmd);

private:
	void DrawBatch();

	// 32-bit needed until the render logic is fixed!
	typedef uint32_t Index;
	const IndexType::Id indexType = IndexType::UInt32;

	uint32_t renderPassIndex;

#pragma pack(push, 1)
	struct Vertex
	{
		vec2_t position;
		vec2_t texCoords;
		uint32_t color;
	};
#pragma pack(pop)
	int maxIndexCount;
	int maxVertexCount;
	int firstIndex;
	int firstVertex;
	int indexCount;
	int vertexCount;
	HRootSignature rootSignature;
	HDescriptorTable descriptorTable;
	HPipeline pipeline;
	HBuffer indexBuffer;
	HBuffer vertexBuffer;
	Index* indices;
	Vertex* vertices;
	uint32_t color;
	const shader_t* shader;
	bool ddhi; // direct descriptor heap indexing
};

struct ImGUI
{
	HTexture Init(bool ddhi, const ShaderByteCode& vs, const ShaderByteCode& ps, TextureFormat::Id rtFormat, HDescriptorTable descTable, RootSignatureDesc* rootSigDesc);
	void RegisterFontAtlas(uint32_t fontIndex);
	void Draw(HTexture renderTarget);
	void BeginFrame();
	void EndFrame();

private:
	struct FrameResources
	{
		HBuffer indexBuffer;
		HBuffer vertexBuffer;
	};

	HRootSignature rootSignature;
	HDescriptorTable descriptorTable;
	HPipeline pipeline;
	HTexture fontAtlas;
	FrameResources frameResources[FrameCount];
	bool frameStarted = false;
	bool ddhi = false;
};

struct Im3D
{
	struct Shader
	{
		enum Id
		{
			PointVS,
			PointPS,
			LineVS,
			LinePS,
			TriangleVS,
			TrianglePS,
			Count
		};
	};

	void Init(bool ddhi_, const ShaderByteCode shaders[Shader::Count], TextureFormat::Id rtFormat);
	void BeginFrame();
	void Draw(const drawSceneViewCommand_t& cmd, HTexture colorTarget, HTexture depthTarget);
	void DrawGUI();
	bool ShouldDrawShapes();
	bool ShouldDrawGizmos();

private:
	HRootSignature rootSignature;
	HDescriptorTable descriptorTables[RHI::FrameCount];
	HPipeline pointPipeline;
	HPipeline linePipeline;
	HPipeline trianglePipeline;
	HBuffer vertexBuffers[RHI::FrameCount];
	HBuffer indexBuffer;
	int prevTime = INT_MIN;
	bool ddhi;
	bool drawGizmos = true;
	bool drawIm3d = true;
};

struct Nuklear
{
	void Init(bool ddhi, const ShaderByteCode& vs, const ShaderByteCode& ps, TextureFormat::Id rtFormat, HDescriptorTable descTable, RootSignatureDesc* rootSigDesc);
	void BeginFrame();
	void Begin(HTexture renderTarget);
	void End();
	void Upload(const nuklearUploadCommand_t& cmd);
	void Draw(const nuklearDrawCommand_t& cmd);

private:
	struct FrameResources
	{
		HBuffer indexBuffer;
		HBuffer vertexBuffer;
	};

	HRootSignature rootSignature;
	HDescriptorTable descriptorTable;
	HPipeline pipeline;
	FrameResources frameResources[FrameCount];
	uint32_t renderPassIndex;
	int prevScissorRect[4];

	// reset every frame
	int firstVertex;
	int firstIndex;
	int numVertexes; // set in Upload
	int numIndexes; // set in Upload
	bool ddhi = false;
};

struct SRP
{
	uint32_t BeginRenderPass(const char* name, float r, float g, float b);
	void EndRenderPass(uint32_t index);

	// @NOTE: SRP::BeginFrame doesn't call RHI::BeginFrame
	// @NOTE: SRP::EndFrame calls RHI::EndFrame and Sys_V_EndFrame
	void BeginFrame(); // call at the start of IRenderPipeline::BeginFrame
	void EndFrame(); // call at the end of IRenderPipeline::EndFrame

	void DrawGUI();

	// call this in Init but only on srp.firstInit
	// you need to register them in your own local descriptor table(s)
	void CreateShaderTraceBuffers();

	bool firstInit = true; // first RP init after a RHI init?
	RenderMode::Id renderMode; // necessary for sampler selection, useful for debugging

	// shader trace
	HBuffer traceRenderBuffer;
	HBuffer traceReadbackBuffer;

	// data for frame breakdown and frame graph
	RenderPassFrame renderPasses[FrameCount];
	RenderPassFrame tempRenderPasses;
	RenderPassStats renderPassStats[MaxRenderPasses];
	RenderPassStats wholeFrameStats;
	FrameStats frameStats;
	bool enableRenderPassQueries = true;

	// PSO stats
	bool psoStatsValid = false;
	int psoCount = 0;
	int psoChangeCount = 0;
};

extern SRP srp;

struct ScopedRenderPass
{
	ScopedRenderPass(const char* name, float r, float g, float b)
	{
		index = srp.BeginRenderPass(name, r, g, b);
	}

	~ScopedRenderPass()
	{
		srp.EndRenderPass(index);
	}

	uint32_t index;
};

struct ScopedDebugLabel
{
	ScopedDebugLabel(const char* name, float r, float g, float b)
	{
		CmdBeginDebugLabel(name, r, g, b);
	}

	~ScopedDebugLabel()
	{
		CmdEndDebugLabel();
	}
};

#define SCOPED_RENDER_PASS(Name, R, G, B) ScopedRenderPass CONCAT(rp_, __LINE__)(Name, R, G, B)
#define SCOPED_DEBUG_LABEL(Name, R, G, B) ScopedDebugLabel CONCAT(rp_, __LINE__)(Name, R, G, B)

#define BASE_SAMPLER_COUNT ((int)(TW_COUNT * TextureFilter::Count * MaxTextureMips))

const image_t* GetBundleImage(const textureBundle_t& bundle);
uint32_t GetBaseSamplerIndex(textureWrap_t wrap, TextureFilter::Id filter, uint32_t minLOD);
uint32_t GetSamplerIndex(textureWrap_t wrap, TextureFilter::Id filter, uint32_t minLOD = 0);
uint32_t GetSamplerIndex(const image_t* image);
void ReadTextureImage(void* outPixels, HTexture hreadback, int w, int h, int alignment, colorSpace_t colorSpace);
void UpdateEntityData(bool& depthHack, int entityNum, double originalTime);
cullType_t GetMirrorredCullType(cullType_t cullType);
uint32_t AlphaTestShaderConstFromStateBits(unsigned int stateBits);

inline void CmdSetViewportAndScissor(uint32_t x, uint32_t y, uint32_t w, uint32_t h)
{
	CmdSetViewport(x, y, w, h);
	CmdSetScissor(x, y, w, h);
}

inline void CmdSetViewportAndScissor(const viewParms_t& vp)
{
	CmdSetViewportAndScissor(vp.viewportX, vp.viewportY, vp.viewportWidth, vp.viewportHeight);
}

inline bool IsDepthFadeEnabled(const shader_t& shader)
{
	return
		r_depthFade->integer != 0 &&
		shader.dfType > DFT_NONE &&
		shader.dfType < DFT_TBD;
}

inline bool IsViewportFullscreen(const viewParms_t& vp)
{
	return
		vp.viewportX == 0 &&
		vp.viewportY == 0 &&
		vp.viewportWidth == glConfig.vidWidth &&
		vp.viewportHeight == glConfig.vidHeight;
}