#include "worldrender.h"
#include "tier1/utlstring.h"
#include "cglm/cglm.h"
#include "cglm/quat.h"
#include "cglm/mat4.h"
#include "game.h"

class CFunnyViewport: public IViewport
{
public:
	virtual void UpdateResolution( uint32_t uX, uint32_t uY ) override;

	virtual IImage *GetRenderImage() override;
	virtual IImage *GetDepthImage() override;
	virtual IImage *GetNormalImage() override;

	virtual void SetFOV( float f ) override;
	virtual void SetPosition( Vector v ) override;
	virtual void SetRotation( Quat q ) override;

	virtual uint32_t GetWidth() override;
	virtual uint32_t GetHeight() override;


	IImage *m_pColor = 0;
	IImage *m_pNormal = 0;
	IImage *m_pWorldCoords = 0;
	IImage *m_pDepth = 0;
	IImage *m_pColorResolved = 0;
	IImage *m_pRenderImage = 0;

	uint32_t m_uX;
	uint32_t m_uY;
	float m_fFov;
	vec3 m_vPos;
	versor m_qRot;
};

void CFunnyViewport::UpdateResolution( uint32_t uX, uint32_t uY )
{
	if (m_pColor)
	{
		g_pRenderContext->DestroyImage(m_pColor);
		g_pRenderContext->DestroyImage(m_pNormal);
		g_pRenderContext->DestroyImage(m_pWorldCoords);
		g_pRenderContext->DestroyImage(m_pDepth);
		g_pRenderContext->DestroyImage(m_pColorResolved);
		g_pRenderContext->DestroyImage(m_pRenderImage);
	}
	m_pColor = g_pRenderContext->CreateRenderTarget(uX, uY, IMAGE_FORMAT_RGBA8_UNORM, MULTISAMPLE_TYPE_4_SAMPLES);
	m_pNormal = g_pRenderContext->CreateRenderTarget(uX, uY, IMAGE_FORMAT_RGBA16_SNORM, MULTISAMPLE_TYPE_4_SAMPLES);
	m_pWorldCoords = g_pRenderContext->CreateRenderTarget(uX, uY, IMAGE_FORMAT_RGBA32_SFLOAT, MULTISAMPLE_TYPE_4_SAMPLES);
	m_pDepth = g_pRenderContext->CreateRenderTarget(uX, uY, IMAGE_FORMAT_D32_SFLOAT, MULTISAMPLE_TYPE_4_SAMPLES);
	m_pColorResolved = g_pRenderContext->CreateRenderTarget(uX, uY, IMAGE_FORMAT_RGBA8_UNORM, MULTISAMPLE_TYPE_1_SAMPLES);
	m_pRenderImage = g_pRenderContext->CreateStorageImage(uX, uY, IMAGE_FORMAT_RGBA16_SFLOAT, MULTISAMPLE_TYPE_1_SAMPLES);
	m_uX = uX;
	m_uY = uY;
}

IImage *CFunnyViewport::GetRenderImage()
{
	return m_pRenderImage;
}

IImage *CFunnyViewport::GetDepthImage()
{
	return m_pDepth;
}

IImage *CFunnyViewport::GetNormalImage()
{
	return m_pNormal;
}

void CFunnyViewport::SetFOV( float f )
{
	m_fFov = f;

}

void CFunnyViewport::SetPosition( Vector v )
{
	m_vPos[0] = v.x;
	m_vPos[1] = v.y;
	m_vPos[2] = v.z;

}

void CFunnyViewport::SetRotation( Quat q )
{
	m_qRot[0] = q.x;
	m_qRot[1] = q.y;
	m_qRot[2] = q.z;
	m_qRot[3] = q.w;

}

uint32_t CFunnyViewport::GetWidth()
{
	return m_uX;
}

uint32_t CFunnyViewport::GetHeight()
{
	return m_uY;
}



struct ViewBuffer_t
{
	mat4 m_matCameraProjection;
	vec4 m_vCameraPosition;
};

struct PerMeshData_t
{
	mat4 m_matTranslation;
	mat4 m_matRotation;
	MaterialData_t m_textureData;
};


class CFunnyMeshInstance;
class CFunnyMesh: public IMesh
{
public:

	virtual void SetVertices( IVertexBuffer *pBuffer ) override;
	virtual void SetIndicies( IIndexBuffer *pBuffer, EIndexFormat eIndexFormat ) override;
	virtual void SetMaterial( IMaterial *pMaterial ) override;


	CUtlString m_szName;
	CUtlVector<CFunnyMeshInstance*> m_instances;
	IMaterial *m_pMaterial;
	IVertexBuffer *m_pVertexBuffer;
};

void CFunnyMesh::SetVertices( IVertexBuffer *pBuffer )
{
	m_pVertexBuffer = pBuffer;
}

void CFunnyMesh::SetIndicies( IIndexBuffer *pBuffer, EIndexFormat eIndexFormat )
{
}

void CFunnyMesh::SetMaterial( IMaterial *pMaterial )
{
	m_pMaterial = pMaterial;
}


class CFunnyMeshInstance: public IMeshInstance
{
public:
	virtual void SetPosition( Vector vPosition ) override;
	virtual void SetRotation( Quat vRotation ) override;
	virtual void SetScale( Vector vScale ) override;
	virtual void SetMaterial( CBaseMaterial *pMaterial ) override;
	virtual void Frame();

	PerMeshData_t m_data = {};
	Quat m_vRotation = { 0, 0, 0, 1 };
	Vector m_vPosition = { 0, 0, 0};
	Vector m_vScale = { 1, 1, 1 };

};

void CFunnyMeshInstance::SetPosition( Vector vPosition )
{	
	m_vPosition = vPosition;
}

void CFunnyMeshInstance::SetRotation( Quat vRotation )
{
	m_vRotation = vRotation;
}

void CFunnyMeshInstance::SetScale( Vector vScale )
{
	m_vScale = vScale;

}
void CFunnyMeshInstance::Frame()
{
	vec3 v; 
	versor q;
	mat4 m;
	glm_mat4_identity(m_data.m_matTranslation);
	glm_mat4_identity(m_data.m_matRotation);
	v[0] = m_vPosition.x;
	v[1] = m_vPosition.y;
	v[2] = m_vPosition.z;
	glm_translate(m_data.m_matTranslation, v);
	q[0] = m_vRotation.x;
	q[1] = m_vRotation.y;
	q[2] = m_vRotation.z;
	q[3] = m_vRotation.w;
	glm_quat_mat4(q, m);
	glm_mat4_mul(m_data.m_matTranslation, m, m_data.m_matTranslation);
	glm_mat4_mul(m_data.m_matRotation, m, m_data.m_matRotation);
	v[0] = m_vScale.x;
	v[1] = m_vScale.y;
	v[2] = m_vScale.z;
	glm_scale_make(m, v);
	glm_mat4_mul(m_data.m_matTranslation, m, m_data.m_matTranslation);
}

void CFunnyMeshInstance::SetMaterial( CBaseMaterial *pMaterial )
{
	pMaterial->SetUpMesh(&m_data.m_textureData);

}


class CFunnyWorldRenderer: public IWorldRenderer
{
public:
	virtual void Init() override;
	virtual void Tick( float fDelta ) override;
	virtual void Frame( float fDelta ) override;
	virtual void Shutdown() override;

	virtual IMesh *CreateMesh( const char *szName ) override;
	virtual IMeshInstance *CreateInstance( IMesh *pMesh ) override;
	virtual void DestroyMeshInstance( IMesh *pMesh, IMeshInstance *pInstance ) override;
	virtual void DestroyMesh( IMesh *pMesh ) override;

	virtual void ConfigureShader( IShader *pShader ) override;

	virtual IViewport *CreateViewport( uint32_t uX, uint32_t uY ) override;
	virtual void DestroyViewport( IViewport *pViewport ) override;

	virtual ITextureArray *GetTextures() override;
private:
	CUtlVector<CFunnyMesh*> m_pMeshes;
	IRenderCommandList *m_pRasterCommandList = NULL;
	IBuffer *m_pViewBuffer;
	ViewBuffer_t *m_pViewBufferData;

	ITextureArray *m_pTextures;
	IShader *m_pRasterShader;
	IMaterial *m_pRasterMaterial;

	IComputeShader *m_pPostProcessing;
	IMaterial *m_pPostProcessingMaterial;

	CUtlVector<CFunnyViewport*> m_viewports;
};

void CFunnyWorldRenderer::ConfigureShader( IShader *pShader )
{
	pShader->AddLayout(0, 32);
	pShader->AddAttribute(0, 0, VERTEX_FORMAT_XYZ32_SFLOAT, 0);
	pShader->AddAttribute(0, 1, VERTEX_FORMAT_XY32_SFLOAT, 12);
	pShader->AddAttribute(0, 2, VERTEX_FORMAT_XYZ32_SFLOAT, 20);
	// albedo
	pShader->AddOutputImage(0, IMAGE_FORMAT_RGBA8_UNORM);
	pShader->AddOutputImage(1, IMAGE_FORMAT_RGBA16_SNORM);
	pShader->AddOutputImage(2, IMAGE_FORMAT_RGBA32_SFLOAT);

	pShader->SetMultisampling(MULTISAMPLE_TYPE_4_SAMPLES);

	pShader->SetDepthImage(IMAGE_FORMAT_D32_SFLOAT);
}


static CFunnyWorldRenderer s_renderer;
IWorldRenderer *g_pWorldRenderer = &s_renderer;

void CFunnyWorldRenderer::Init()
{
	m_pRasterShader = g_pRenderContext->CreateShader("game/core/shaders/mesh_raster.shader_c");
	ConfigureShader(m_pRasterShader);
	m_pRasterShader->Build();
	m_pRasterMaterial = g_pRenderContext->CreateMaterial(m_pRasterShader);

	m_pRasterCommandList = g_pRenderContext->CreateCommandList();
	m_pViewBuffer = g_pRenderContext->CreateConstantBuffer(sizeof(ViewBuffer_t));
	m_pTextures = g_pRenderContext->CreateTextureArray();

	IRayTracingShader *prt = g_pRenderContext->CreateRayShader("game/core/shaders/mesh_trace.shader_c");
	prt->AddShader("error", "game/core/shaders/funny_error.shader_c");
	prt->Build();

	m_pPostProcessing = g_pRenderContext->CreateComputeShader("game/core/shaders/agx.shader_c");
	m_pPostProcessing->Build();

	m_pPostProcessingMaterial = g_pRenderContext->CreateMaterial(m_pPostProcessing);
}

void CFunnyWorldRenderer::Tick( float fDelta )
{
}

void CFunnyWorldRenderer::Frame( float fDelta )
{
	for (auto &v: m_viewports)
	{
		uint32_t uWidth = v->m_uX;
		uint32_t uHeight = v->m_uY;
		mat4 matCamera;
		mat4 matCamera2;
		glm_mat4_identity(matCamera);
		glm_mat4_identity(matCamera2);
		glm_translate(matCamera2, v->m_vPos);
		glm_quat_rotate(matCamera2, v->m_qRot, matCamera2);
		glm_mat4_inv(matCamera2, matCamera2);
		glm_perspective(glm_rad(v->m_fFov), uWidth/(float)uHeight, 0.01, 10000, matCamera);
		glm_mul(matCamera, matCamera2, matCamera);
		m_pViewBufferData = (ViewBuffer_t*)m_pViewBuffer->Map();
		m_pViewBuffer->Lock();
		V_memcpy(&m_pViewBufferData->m_matCameraProjection, matCamera, sizeof(matCamera));
		V_memcpy(&m_pViewBufferData->m_vCameraPosition, v->m_vPos, sizeof(vec3));
		m_pViewBuffer->Unlock();
		m_pViewBuffer->Unmap();

		m_pRasterCommandList->StartRecording();
		m_pRasterCommandList->SetRenderResolution(uWidth, uHeight);
		m_pRasterCommandList->SetRenderTarget(0, v->m_pColor);
		m_pRasterCommandList->SetRenderTarget(1, v->m_pNormal);
		m_pRasterCommandList->SetRenderTarget(2, v->m_pWorldCoords);
		m_pRasterCommandList->SetDepthTarget(v->m_pDepth);
		m_pRasterCommandList->SetViewport(0, 0, uWidth, uHeight, 0, 1);
		m_pRasterCommandList->SetScissors(0, 0, uWidth, uHeight);
		m_pRasterCommandList->SetClearColor(0, 0, 0, 0, 0);
		m_pRasterCommandList->SetClearColor(1, 0, 0, 0, 0);
		m_pRasterCommandList->SetClearColor(2, 0, 0, 0, 0);
		m_pRasterCommandList->SetClearDepth(1);
		m_pRasterCommandList->SetLoadStoreModes(0, LOAD_MODE_CLEAR, STORE_MODE_STORE);
		m_pRasterCommandList->SetLoadStoreModes(1, LOAD_MODE_CLEAR, STORE_MODE_STORE);
		m_pRasterCommandList->SetLoadStoreModes(2, LOAD_MODE_CLEAR, STORE_MODE_STORE);
		uint32_t uTotalMeshes = 0;
		uint32_t u = 0;
		for ( auto mesh: m_pMeshes)
		{
			uTotalMeshes+=mesh->m_instances.GetSize();
		}
		if (!uTotalMeshes)
			return;
		IBuffer *pDataBuffer = g_pRenderContext->CreateStorageBuffer(uTotalMeshes*sizeof(PerMeshData_t));
		pDataBuffer->Lock();
		void *pData = pDataBuffer->Map();
		for ( auto mesh: m_pMeshes)
		{
			if (mesh->m_instances.GetSize()==0)
				continue;
			for ( uint32_t i = 0; i < mesh->m_instances.GetSize(); i++ )
			{
				mesh->m_instances[i]->Frame();
				V_memcpy(&((PerMeshData_t*)pData)[i+u], &mesh->m_instances[i]->m_data, sizeof(PerMeshData_t));
			}
			u+=mesh->m_instances.GetSize();
		}
		pDataBuffer->Unmap();
		pDataBuffer->Unlock();
		g_pRenderContext->DestroyBuffer(pDataBuffer);

		m_pRasterMaterial->VSSetConstantsBuffer(0, m_pViewBuffer);
		m_pRasterMaterial->VSSetConstantsBuffer(1, pDataBuffer);
		m_pRasterMaterial->PSSetTextureArray(1, m_pTextures);

		m_pRasterCommandList->SetMaterial(m_pRasterMaterial);
		u = 0;
		for ( auto mesh: m_pMeshes)
		{
			if (mesh->m_instances.GetSize()==0)
				continue;
			m_pRasterCommandList->SetVertexBuffer(0, mesh->m_pVertexBuffer);
			m_pRasterCommandList->DrawPrimitives(mesh->m_pVertexBuffer->GetSize()/32, 0, mesh->m_instances.GetSize(), u);
			u += mesh->m_instances.GetSize();
		}

		m_pRasterCommandList->ResolveImage(v->m_pColor, v->m_pColorResolved);

		m_pPostProcessingMaterial->CSSetTexture(0, 0, v->m_pColorResolved);
		m_pPostProcessingMaterial->CSSetTexture(1, 0, v->m_pRenderImage);
		m_pRasterCommandList->SetMaterial(m_pPostProcessingMaterial);
		m_pRasterCommandList->Barrier(v->m_pColorResolved, true, false);
		m_pRasterCommandList->Barrier(v->m_pRenderImage, false, true);
		m_pRasterCommandList->DispatchCompute(uWidth, uHeight, 1);
		m_pRasterCommandList->EndRecording();

		g_pRenderContext->SubmitCommandList(m_pRasterCommandList);
	}
}

void CFunnyWorldRenderer::Shutdown()
{

}

IMesh *CFunnyWorldRenderer::CreateMesh( const char *szName )
{
	CFunnyMesh *pMesh = new CFunnyMesh;
	pMesh->m_szName = szName;
	m_pMeshes.AppendTail(pMesh);
	return pMesh;
}

IMeshInstance *CFunnyWorldRenderer::CreateInstance( IMesh *pMesh )
{
	CFunnyMesh *pFunnyMesh = (CFunnyMesh*)pMesh;
	CFunnyMeshInstance *pInstance = new CFunnyMeshInstance;
	pFunnyMesh->m_instances.AppendTail(pInstance);
	return pInstance;
}

void CFunnyWorldRenderer::DestroyMeshInstance( IMesh *pMesh, IMeshInstance *pInstance )
{
	if (!pMesh)
		return;
	CFunnyMesh *pFunnyMesh = (CFunnyMesh*)pMesh;
	for ( uint32_t i = 0; i < pFunnyMesh->m_instances.GetSize(); i++ )
	{
		if ( pFunnyMesh->m_instances[i] != pInstance)
			continue;
		delete pFunnyMesh->m_instances[i];
		pFunnyMesh->m_instances.RemoveAt(i, 1);
		break;

	}

}

void CFunnyWorldRenderer::DestroyMesh( IMesh *pMesh )
{

}


ITextureArray *CFunnyWorldRenderer::GetTextures()
{
	return m_pTextures;

}

IViewport *CFunnyWorldRenderer::CreateViewport( uint32_t uX, uint32_t uY )
{
	CFunnyViewport *pViewport = new CFunnyViewport;
	pViewport->UpdateResolution(uX, uY);
	m_viewports.AppendTail(pViewport);
	return pViewport;
}

void CFunnyWorldRenderer::DestroyViewport( IViewport *pViewport )
{
	for ( uint32_t i = 0; i < m_viewports.GetSize(); i++ )
	{
		if (m_viewports[i] == pViewport)
		{
			m_viewports.RemoveAt(i);
			delete (CFunnyViewport*)pViewport;
			return;
		}
	}
}