/*****************************************************************************
					The Dark Mod GPL Source Code

 This file is part of the The Dark Mod Source Code, originally based
 on the Doom 3 GPL Source Code as published in 2011.

 The Dark Mod Source Code 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 3 of the License,
 or (at your option) any later version. For details, see LICENSE.TXT.

 Project: The Dark Mod (http://www.thedarkmod.com/)

******************************************************************************/

#include "precompiled.h"

#pragma hdrstop

#include "AmbientOcclusionStage.h"
#include "Image.h"
#include "tr_local.h"
#include "GLSLProgramManager.h"
#include "GLSLProgram.h"
#include "Profiling.h"
#include "GLSLUniforms.h"
#include "glsl.h"
#include "FrameBufferManager.h"
#include "FrameBuffer.h"

idCVar r_ssao("r_ssao", "0", CVAR_INTEGER | CVAR_RENDERER | CVAR_ARCHIVE, "Screen space ambient occlusion: 0 - off, 1 - low, 2 - medium, 3 - high");
idCVar r_ssao_radius("r_ssao_radius", "32", CVAR_FLOAT | CVAR_RENDERER | CVAR_ARCHIVE,
	"View space sample radius - larger values provide a softer, spread effect, but risk causing unwanted halo shadows around objects");
idCVar r_ssao_bias("r_ssao_bias", "0.05", CVAR_FLOAT | CVAR_RENDERER | CVAR_ARCHIVE,
	"Min depth difference to count for occlusion, used to avoid some acne effects");
idCVar r_ssao_intensity("r_ssao_intensity", "1.0", CVAR_FLOAT | CVAR_RENDERER | CVAR_ARCHIVE,
	"SSAO intensity factor, the higher the value, the stronger the effect");
idCVar r_ssao_base("r_ssao_base", "0.1", CVAR_FLOAT | CVAR_RENDERER | CVAR_ARCHIVE,
	"Minimum baseline visibility below which AO cannot drop");
idCVar r_ssao_edgesharpness("r_ssao_edgesharpness", "1", CVAR_FLOAT | CVAR_RENDERER | CVAR_ARCHIVE, "Edge sharpness in SSAO blur");

extern idCVar r_fboResolution;

AmbientOcclusionStage ambientOcclusionImpl;
AmbientOcclusionStage *ambientOcclusion = &ambientOcclusionImpl;

namespace {
	struct AOUniforms : GLSLUniformGroup {
		UNIFORM_GROUP_DEF(AOUniforms)

		DEFINE_UNIFORM(sampler, depthTexture)
		DEFINE_UNIFORM(float, sampleRadius)
		DEFINE_UNIFORM(float, depthBias)
		DEFINE_UNIFORM(float, baseValue)
		DEFINE_UNIFORM(int, numSamples)
		DEFINE_UNIFORM(int, numSpiralTurns)
		DEFINE_UNIFORM(float, intensityDivR6)
		DEFINE_UNIFORM(int, maxMipLevel)
	};

	struct BlurUniforms : GLSLUniformGroup {
		UNIFORM_GROUP_DEF(BlurUniforms)

		DEFINE_UNIFORM(sampler, source)
		DEFINE_UNIFORM(vec2, axis)
		DEFINE_UNIFORM(float, edgeSharpness)
	};

	struct DepthMipUniforms : GLSLUniformGroup {
		UNIFORM_GROUP_DEF(DepthMipUniforms)

		DEFINE_UNIFORM(sampler, depth)
		DEFINE_UNIFORM(int, previousMipLevel)
	};

	void RenderTexture(idImage *image) {
		image->type = TT_2D;
	}

	void LoadSSAOShader(GLSLProgram *ssaoShader) {
		ssaoShader->LoadFromFiles( "ssao.vert.glsl", "ssao.frag.glsl" );
		AOUniforms *uniforms = ssaoShader->GetUniformGroup<AOUniforms>();
		uniforms->depthTexture.Set(0);
	}

	void LoadSSAOBlurShader(GLSLProgram *blurShader) {
		blurShader->LoadFromFiles( "ssao.vert.glsl", "ssao_blur.frag.glsl" );
		BlurUniforms *uniforms = blurShader->GetUniformGroup<BlurUniforms>();
		uniforms->source.Set(0);
	}

	void CreateSSAOColorFBO(FrameBuffer *fbo, idImage *color) {
		fbo->Init( frameBuffers->renderWidth, frameBuffers->renderHeight );
		color->GenerateAttachment( frameBuffers->renderWidth, frameBuffers->renderHeight, GL_RGBA8, GL_LINEAR, GL_CLAMP_TO_EDGE );
		fbo->AddColorRenderTexture( 0, color );
	}

	void CreateViewspaceDepthFBO(FrameBuffer *fbo, idImage *image, int mipLevel) {
		// create texture, if necessary
		image->type = TT_2D;
		if (image->texnum == idImage::TEXTURE_NOT_LOADED || image->uploadWidth != frameBuffers->renderWidth || image->uploadHeight != frameBuffers->renderHeight ) {
			image->PurgeImage();
			qglGenTextures(1, &image->texnum);
			image->uploadWidth = frameBuffers->renderWidth;
			image->uploadHeight = frameBuffers->renderHeight;
			qglBindTexture(GL_TEXTURE_2D, image->texnum);
			for (int i = 0; i <= AmbientOcclusionStage::MAX_DEPTH_MIPS; ++i) {
				qglTexImage2D(GL_TEXTURE_2D, i, GL_R32F, image->uploadWidth >> i, image->uploadHeight >> i, 0, GL_RED, GL_FLOAT, nullptr);
			}
			qglTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_BASE_LEVEL, 0);
			qglTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAX_LEVEL, AmbientOcclusionStage::MAX_DEPTH_MIPS);
			qglTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
			qglTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
			qglTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_MIRRORED_REPEAT);
			qglTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_MIRRORED_REPEAT);
			GL_SetDebugLabel( GL_TEXTURE, image->texnum, image->imgName );
		}

		fbo->Init( frameBuffers->renderWidth >> mipLevel, frameBuffers->renderHeight >> mipLevel );
		fbo->AddColorRenderTexture( 0, image, mipLevel );
	}
}

void AmbientOcclusionStage::Init() {
	ssaoResult = globalImages->ImageFromFunction("SSAO ColorBuffer", RenderTexture);
	ssaoBlurred = globalImages->ImageFromFunction("SSAO Blurred", RenderTexture);
	viewspaceDepth = globalImages->ImageFromFunction("SSAO Depth", RenderTexture);

	ssaoFBO = frameBuffers->CreateFromGenerator( "ssao_color", [this](FrameBuffer *fbo) { CreateSSAOColorFBO( fbo, ssaoResult ); } );
	ssaoBlurFBO = frameBuffers->CreateFromGenerator( "ssao_blurred", [this](FrameBuffer *fbo) { CreateSSAOColorFBO( fbo, ssaoBlurred ); } );
	for (int i = 0; i <= MAX_DEPTH_MIPS; ++i) {
		depthMipFBOs[i] = frameBuffers->CreateFromGenerator( idStr("ssao_depth_") + i, [this, i](FrameBuffer *fbo) { CreateViewspaceDepthFBO( fbo, viewspaceDepth, i ); } );
	}
	frameBuffers->defaultFbo->Bind();

	ssaoShader = programManager->LoadFromGenerator("ssao", LoadSSAOShader);
	ssaoBlurShader = programManager->LoadFromGenerator("ssao_blur", LoadSSAOBlurShader);
	depthShader = programManager->LoadFromFiles("ssao_depth", "ssao.vert.glsl", "ssao_depth.frag.glsl");
	depthMipShader = programManager->LoadFromFiles("ssao_depth_mip", "ssao.vert.glsl", "ssao_depthmip.frag.glsl");
	showSSAOShader = programManager->LoadFromFiles("ssao_show", "ssao.vert.glsl", "ssao_show.frag.glsl");
}

void AmbientOcclusionStage::Shutdown() {
	if (ssaoFBO != nullptr) {
		ssaoFBO->Destroy();
		ssaoFBO = nullptr;
	}
	if (ssaoBlurFBO != nullptr) {
		ssaoBlurFBO->Destroy();
		ssaoBlurFBO = nullptr;
	}
	for (int i = 0; i <= MAX_DEPTH_MIPS; ++i) {
		if (depthMipFBOs[i] != nullptr) {
			depthMipFBOs[i]->Destroy();
			depthMipFBOs[i] = nullptr;
		}
	}
	if (viewspaceDepth != nullptr) {
		viewspaceDepth->PurgeImage();
		viewspaceDepth = nullptr;
	}
	if (ssaoResult != nullptr) {
		ssaoResult->PurgeImage();
		ssaoResult = nullptr;
	}
	if (ssaoBlurred != nullptr) {
		ssaoBlurred->PurgeImage();
		ssaoBlurred = nullptr;
	}
}

extern GLuint fboPrimary;
extern bool primaryOn;

void AmbientOcclusionStage::ComputeSSAOFromDepth() {
	GL_PROFILE("AmbientOcclusionStage");

	if (ssaoFBO == nullptr) {
		Init();
	}

	PrepareDepthPass();
	SSAOPass();
	BlurPass();

	frameBuffers->currentRenderFbo->Bind();
}

void AmbientOcclusionStage::SSAOPass() {
	GL_PROFILE("SSAOPass");

	ssaoFBO->Bind();
	qglClearColor(1, 1, 1, 1);
	qglClear(GL_COLOR_BUFFER_BIT);
	GL_SelectTexture(0);
	viewspaceDepth->Bind();

	ssaoShader->Activate();
	SetQualityLevelUniforms();
	RB_DrawFullScreenQuad();
}

void AmbientOcclusionStage::BlurPass() {
	GL_PROFILE("BlurPass");

	ssaoBlurShader->Activate();
	BlurUniforms *uniforms = ssaoBlurShader->GetUniformGroup<BlurUniforms>();
	uniforms->edgeSharpness.Set(r_ssao_edgesharpness.GetFloat());
	uniforms->source.Set(0);
	uniforms->axis.Set(1, 0);

	// first horizontal pass
	ssaoBlurFBO->Bind();
	qglClearColor(1, 1, 1, 1);
	qglClear(GL_COLOR_BUFFER_BIT);
	GL_SelectTexture(0);
	ssaoResult->Bind();
	RB_DrawFullScreenQuad();

	// second vertical pass
	uniforms->axis.Set(0, 1);
	ssaoFBO->Bind();
	qglClear(GL_COLOR_BUFFER_BIT);
	ssaoBlurred->Bind();
	RB_DrawFullScreenQuad();
}

void AmbientOcclusionStage::BindSSAOTexture(int index) {
	GL_SelectTexture(index);
	if (ShouldEnableForCurrentView()) {
		ssaoResult->Bind();
	}
	else {
		globalImages->whiteImage->Bind();
	}
}

bool AmbientOcclusionStage::ShouldEnableForCurrentView() const {
	return r_ssao.GetBool() && !backEnd.viewDef->IsLightGem() && !backEnd.viewDef->isSubview && !backEnd.viewDef->renderWorld->mapName.IsEmpty();
}

void AmbientOcclusionStage::PrepareDepthPass() {
	GL_PROFILE("PrepareDepthPass");

	depthMipFBOs[0]->Bind();
	GL_ScissorRelative( 0, 0, 1, 1 );
	qglClear(GL_COLOR_BUFFER_BIT);
	GL_SelectTexture(0);
	globalImages->currentDepthImage->Bind();

	depthShader->Activate();
	RB_DrawFullScreenQuad();

	if (r_ssao.GetInteger() > 1) {
		GL_PROFILE("DepthMips");
		// generate mip levels - used by the AO shader for distant samples to ensure we hit the texture cache as much as possible
		depthMipShader->Activate();
		DepthMipUniforms *uniforms = depthMipShader->GetUniformGroup<DepthMipUniforms>();
		uniforms->depth.Set(0);
		viewspaceDepth->Bind();
		for (int i = 1; i <= MAX_DEPTH_MIPS; ++i) {
			depthMipFBOs[i]->Bind();
			qglClear(GL_COLOR_BUFFER_BIT);
			uniforms->previousMipLevel.Set(i - 1);
			RB_DrawFullScreenQuad();
		}
	}
}

void AmbientOcclusionStage::ShowSSAO() {
	showSSAOShader->Activate();
	BindSSAOTexture(0);
	RB_DrawFullScreenQuad();
}

void AmbientOcclusionStage::SetQualityLevelUniforms() {
	AOUniforms *uniforms = ssaoShader->GetUniformGroup<AOUniforms>();
	uniforms->depthBias.Set(r_ssao_bias.GetFloat());
	uniforms->baseValue.Set(r_ssao_base.GetFloat());
	float sampleRadius;
	switch (r_ssao.GetInteger()) {
	case 1:
		 uniforms->maxMipLevel.Set(0);
		 uniforms->numSpiralTurns.Set(5);
		 uniforms->numSamples.Set(7);
		 sampleRadius = 0.5f * r_ssao_radius.GetFloat();
		 break;
	case 2:
		 uniforms->maxMipLevel.Set(MAX_DEPTH_MIPS);
		 uniforms->numSpiralTurns.Set(7);
		 uniforms->numSamples.Set(12);
		 sampleRadius = r_ssao_radius.GetFloat();
		 break;
	case 3:
	default:
		 uniforms->maxMipLevel.Set(MAX_DEPTH_MIPS);
		 uniforms->numSpiralTurns.Set(7);
		 uniforms->numSamples.Set(24);
		 sampleRadius = 2.0f * r_ssao_radius.GetFloat();
		 break;
	}
	uniforms->intensityDivR6.Set(r_ssao_intensity.GetFloat() / pow(sampleRadius * 0.02309f, 6));
	uniforms->sampleRadius.Set(sampleRadius);
}
