#include "SSAORenderer.hpp"

#define GLFW_INCLUDE_NONE
#include <GLFW/glfw3.h>
#ifdef __APPLE__
#include <OpenGL/gl3.h>
#endif
#include <glm/gtc/type_ptr.hpp>
#include <iostream>
#include <random>

namespace renderer {

SSAORenderer::SSAORenderer ()
    : mWidth( 0 ), mHeight( 0 ), mHalfRes( true ), mKernelSize( 32 ), mGBufferFBO( 0 ), mGNormal( 0 ), mGDepth( 0 ),
      mSSAOFBO( 0 ), mSSAOTexture( 0 ), mSSAOBlurFBO( 0 ), mSSAOBlurTexture( 0 ), mNoiseTexture( 0 ), mQuadVAO( 0 ),
      mQuadVBO( 0 ) {
}

SSAORenderer::~SSAORenderer () {
    if( mGNormal )
        glDeleteTextures( 1, &mGNormal );
    if( mGDepth )
        glDeleteTextures( 1, &mGDepth );
    if( mGBufferFBO )
        glDeleteFramebuffers( 1, &mGBufferFBO );
    if( mSSAOTexture )
        glDeleteTextures( 1, &mSSAOTexture );
    if( mSSAOBlurTexture )
        glDeleteTextures( 1, &mSSAOBlurTexture );
    if( mSSAOFBO )
        glDeleteFramebuffers( 1, &mSSAOFBO );
    if( mSSAOBlurFBO )
        glDeleteFramebuffers( 1, &mSSAOBlurFBO );
    if( mNoiseTexture )
        glDeleteTextures( 1, &mNoiseTexture );
    if( mQuadVBO )
        glDeleteBuffers( 1, &mQuadVBO );
    if( mQuadVAO )
        glDeleteVertexArrays( 1, &mQuadVAO );
}

bool SSAORenderer::init ( int width, int height, bool halfRes ) {
    mWidth   = width;
    mHeight  = height;
    mHalfRes = halfRes;

    if( !mGBufferShader.fromFiles( "apps/openmb/resources/shaders/gbuffer.vert",
                                   "apps/openmb/resources/shaders/gbuffer.frag" ) ) {
        std::cerr << "Failed to load gbuffer shader" << std::endl;
        return false;
    }
    if( !mSSAOShader.fromFiles( "apps/openmb/resources/shaders/ssao.vert",
                                "apps/openmb/resources/shaders/ssao.frag" ) ) {
        std::cerr << "Failed to load ssao shader" << std::endl;
        return false;
    }
    if( !mBlurShader.fromFiles( "apps/openmb/resources/shaders/ssao.vert",
                                "apps/openmb/resources/shaders/ssao_blur.frag" ) ) {
        std::cerr << "Failed to load ssao blur shader" << std::endl;
        return false;
    }

    initBuffers();
    initKernelAndNoise();

    float quadVertices[] = {

        -1.0f,
        -1.0f,
        0.0f,
        0.0f,
        1.0f,
        -1.0f,
        1.0f,
        0.0f,
        -1.0f,
        1.0f,
        0.0f,
        1.0f,
        -1.0f,
        1.0f,
        0.0f,
        1.0f,
        1.0f,
        -1.0f,
        1.0f,
        0.0f,
        1.0f,
        1.0f,
        1.0f,
        1.0f,
    };

    glGenVertexArrays( 1, &mQuadVAO );
    glGenBuffers( 1, &mQuadVBO );
    glBindVertexArray( mQuadVAO );
    glBindBuffer( GL_ARRAY_BUFFER, mQuadVBO );
    glBufferData( GL_ARRAY_BUFFER, sizeof( quadVertices ), quadVertices, GL_STATIC_DRAW );
    glEnableVertexAttribArray( 0 );
    glVertexAttribPointer( 0, 2, GL_FLOAT, GL_FALSE, 4 * sizeof( float ), (void*)0 );
    glEnableVertexAttribArray( 1 );
    glVertexAttribPointer( 1, 2, GL_FLOAT, GL_FALSE, 4 * sizeof( float ), (void*)( 2 * sizeof( float ) ) );
    glBindVertexArray( 0 );

    return true;
}

void SSAORenderer::initBuffers () {

    if( mGBufferFBO == 0 )
        glGenFramebuffers( 1, &mGBufferFBO );
    glBindFramebuffer( GL_FRAMEBUFFER, mGBufferFBO );

    if( mGNormal )
        glDeleteTextures( 1, &mGNormal );
    glGenTextures( 1, &mGNormal );
    glBindTexture( GL_TEXTURE_2D, mGNormal );
    glTexImage2D( GL_TEXTURE_2D, 0, GL_RGB16F, mWidth, mHeight, 0, GL_RGB, GL_FLOAT, nullptr );
    glTexParameteri( GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST );
    glTexParameteri( GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST );
    glFramebufferTexture2D( GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, mGNormal, 0 );

    if( mGDepth )
        glDeleteTextures( 1, &mGDepth );
    glGenTextures( 1, &mGDepth );
    glBindTexture( GL_TEXTURE_2D, mGDepth );
    glTexImage2D( GL_TEXTURE_2D, 0, GL_DEPTH_COMPONENT, mWidth, mHeight, 0, GL_DEPTH_COMPONENT, GL_FLOAT, nullptr );
    glTexParameteri( GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST );
    glTexParameteri( GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST );
    glTexParameteri( GL_TEXTURE_2D, GL_TEXTURE_COMPARE_MODE, GL_NONE );
    glFramebufferTexture2D( GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, GL_TEXTURE_2D, mGDepth, 0 );

    GLenum attachments[1] = { GL_COLOR_ATTACHMENT0 };
    glDrawBuffers( 1, attachments );

    if( glCheckFramebufferStatus( GL_FRAMEBUFFER ) != GL_FRAMEBUFFER_COMPLETE )
        std::cerr << "GBuffer not complete" << std::endl;
    glBindFramebuffer( GL_FRAMEBUFFER, 0 );

    if( mSSAOFBO == 0 )
        glGenFramebuffers( 1, &mSSAOFBO );
    glBindFramebuffer( GL_FRAMEBUFFER, mSSAOFBO );
    if( mSSAOTexture )
        glDeleteTextures( 1, &mSSAOTexture );
    glGenTextures( 1, &mSSAOTexture );
    glBindTexture( GL_TEXTURE_2D, mSSAOTexture );
    int w = mHalfRes ? std::max( 1, mWidth / 2 ) : mWidth;
    int h = mHalfRes ? std::max( 1, mHeight / 2 ) : mHeight;
    glTexImage2D( GL_TEXTURE_2D, 0, GL_R8, w, h, 0, GL_RED, GL_UNSIGNED_BYTE, nullptr );
    glTexParameteri( GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST );
    glTexParameteri( GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST );
    glFramebufferTexture2D( GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, mSSAOTexture, 0 );
    if( glCheckFramebufferStatus( GL_FRAMEBUFFER ) != GL_FRAMEBUFFER_COMPLETE )
        std::cerr << "SSAO FBO not complete" << std::endl;
    glBindFramebuffer( GL_FRAMEBUFFER, 0 );

    if( mSSAOBlurFBO == 0 )
        glGenFramebuffers( 1, &mSSAOBlurFBO );
    glBindFramebuffer( GL_FRAMEBUFFER, mSSAOBlurFBO );
    if( mSSAOBlurTexture )
        glDeleteTextures( 1, &mSSAOBlurTexture );
    glGenTextures( 1, &mSSAOBlurTexture );
    glBindTexture( GL_TEXTURE_2D, mSSAOBlurTexture );
    glTexImage2D( GL_TEXTURE_2D, 0, GL_R8, w, h, 0, GL_RED, GL_UNSIGNED_BYTE, nullptr );
    glTexParameteri( GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST );
    glTexParameteri( GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST );
    glFramebufferTexture2D( GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, mSSAOBlurTexture, 0 );
    if( glCheckFramebufferStatus( GL_FRAMEBUFFER ) != GL_FRAMEBUFFER_COMPLETE )
        std::cerr << "SSAO Blur FBO not complete" << std::endl;
    glBindFramebuffer( GL_FRAMEBUFFER, 0 );
}

void SSAORenderer::initKernelAndNoise () {
    mKernel.clear();
    std::uniform_real_distribution<float> randomFloats( 0.0f, 1.0f );
    std::mt19937                          gen;

    for( int i = 0; i < mKernelSize; ++i ) {
        glm::vec3 sample( randomFloats( gen ) * 2.0f - 1.0f, randomFloats( gen ) * 2.0f - 1.0f,
                          randomFloats( gen ) );
        sample = glm::normalize( sample );
        sample *= randomFloats( gen );
        float scale = float( i ) / float( mKernelSize );
        scale       = glm::mix( 0.1f, 1.0f, scale * scale );
        sample *= scale;
        mKernel.push_back( sample );
    }

    mSSAOShader.use();
    for( int i = 0; i < mKernelSize; ++i ) {
        mSSAOShader.setVec3( std::string( "samples[" ) + std::to_string( i ) + std::string( "]" ), mKernel[i] );
    }
    mSSAOShader.setInt( "kernelSize", mKernelSize );

    std::vector<float> noise;
    for( int i = 0; i < 16; ++i ) {
        float x = randomFloats( gen ) * 2.0f - 1.0f;
        float y = randomFloats( gen ) * 2.0f - 1.0f;
        noise.push_back( x );
        noise.push_back( y );
        noise.push_back( 0.0f );
    }
    if( mNoiseTexture )
        glDeleteTextures( 1, &mNoiseTexture );
    glGenTextures( 1, &mNoiseTexture );
    glBindTexture( GL_TEXTURE_2D, mNoiseTexture );
    glTexImage2D( GL_TEXTURE_2D, 0, GL_RGB32F, 4, 4, 0, GL_RGB, GL_FLOAT, noise.data() );
    glTexParameteri( GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST );
    glTexParameteri( GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST );
    glTexParameteri( GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT );
    glTexParameteri( GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT );
    glBindTexture( GL_TEXTURE_2D, 0 );
}

void SSAORenderer::resize ( int width, int height ) {
    if( width == mWidth && height == mHeight )
        return;
    mWidth  = width;
    mHeight = height;
    initBuffers();
}

void SSAORenderer::bindGBuffer () {
    glBindFramebuffer( GL_FRAMEBUFFER, mGBufferFBO );
    glViewport( 0, 0, mWidth, mHeight );
    glClear( GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT );
}

void SSAORenderer::unbind () {
    glBindFramebuffer( GL_FRAMEBUFFER, 0 );
}

void SSAORenderer::computeSSAO ( const glm::mat4& proj, const glm::mat4& invProj, float radius, float bias,
                                 float power ) {
    int w = mHalfRes ? std::max( 1, mWidth / 2 ) : mWidth;
    int h = mHalfRes ? std::max( 1, mHeight / 2 ) : mHeight;

    glBindFramebuffer( GL_FRAMEBUFFER, mSSAOFBO );
    glViewport( 0, 0, w, h );
    glClear( GL_COLOR_BUFFER_BIT );

    mSSAOShader.use();
    mSSAOShader.setMat4( "proj", proj );
    mSSAOShader.setMat4( "invProj", invProj );
    mSSAOShader.setFloat( "radius", radius );
    mSSAOShader.setFloat( "bias", bias );
    mSSAOShader.setFloat( "power", power );

    glActiveTexture( GL_TEXTURE0 );
    glBindTexture( GL_TEXTURE_2D, mGDepth );
    mSSAOShader.setInt( "gDepth", 0 );

    glActiveTexture( GL_TEXTURE1 );
    glBindTexture( GL_TEXTURE_2D, mGNormal );
    mSSAOShader.setInt( "gNormal", 1 );

    glActiveTexture( GL_TEXTURE2 );
    glBindTexture( GL_TEXTURE_2D, mNoiseTexture );
    mSSAOShader.setInt( "texNoise", 2 );

    float noiseScaleVal = float( w ) / 4.0f;
    mSSAOShader.setInt( "gDepth", 0 );
    mSSAOShader.setInt( "gNormal", 1 );
    mSSAOShader.setInt( "texNoise", 2 );
    mSSAOShader.setFloat( "noiseScale", noiseScaleVal );
    glBindVertexArray( mQuadVAO );
    glDrawArrays( GL_TRIANGLES, 0, 6 );
    glBindVertexArray( 0 );

    glBindFramebuffer( GL_FRAMEBUFFER, 0 );
}

void SSAORenderer::blurSSAO () {
    int w = mHalfRes ? std::max( 1, mWidth / 2 ) : mWidth;
    int h = mHalfRes ? std::max( 1, mHeight / 2 ) : mHeight;

    glBindFramebuffer( GL_FRAMEBUFFER, mSSAOBlurFBO );
    glViewport( 0, 0, w, h );
    glClear( GL_COLOR_BUFFER_BIT );

    mBlurShader.use();
    glActiveTexture( GL_TEXTURE0 );
    glBindTexture( GL_TEXTURE_2D, mSSAOTexture );
    mBlurShader.setInt( "ssaoInput", 0 );
    mBlurShader.setInt( "horizontal", 1 );
    mBlurShader.setVec2( "texelSize", glm::vec2( 1.0f / float( w ), 1.0f / float( h ) ) );
    glBindVertexArray( mQuadVAO );
    glDrawArrays( GL_TRIANGLES, 0, 6 );
    glBindVertexArray( 0 );

    glBindFramebuffer( GL_FRAMEBUFFER, mSSAOFBO );
    glViewport( 0, 0, w, h );
    glClear( GL_COLOR_BUFFER_BIT );

    mBlurShader.use();
    glActiveTexture( GL_TEXTURE0 );
    glBindTexture( GL_TEXTURE_2D, mSSAOBlurTexture );
    mBlurShader.setInt( "ssaoInput", 0 );
    mBlurShader.setInt( "horizontal", 0 );
    mBlurShader.setVec2( "texelSize", glm::vec2( 1.0f / float( w ), 1.0f / float( h ) ) );
    glBindVertexArray( mQuadVAO );
    glDrawArrays( GL_TRIANGLES, 0, 6 );
    glBindVertexArray( 0 );

    glBindFramebuffer( GL_FRAMEBUFFER, 0 );
}

unsigned int SSAORenderer::getSSAOTextureID () const {
    return mSSAOTexture;
}

void SSAORenderer::bindSSAOTextureToUnit ( int unit, renderer::Shader& shader, const std::string& uniformName ) const {
    glActiveTexture( GL_TEXTURE0 + unit );
    glBindTexture( GL_TEXTURE_2D, getSSAOTextureID() );
    shader.setInt( uniformName, unit );
}

}   // namespace renderer