// This file is part of the OGRE project.
// It is subject to the license terms in the LICENSE file found in the top-level directory
// of this distribution and at https://www.ogre3d.org/licensing.
// SPDX-License-Identifier: MIT
#include "OgreTerrainRTShaderSRS.h"

#include "OgrePass.h"
#include "OgreShaderProgramSet.h"
#include "OgreShaderParameter.h"
#include "OgreMaterialSerializer.h"
#include "OgreShaderGenerator.h"
#include "OgreShaderFunction.h"
#include "OgreShaderProgram.h"
#include "OgreTerrainQuadTreeNode.h"
#include "OgreTextureManager.h"

namespace Ogre {
using namespace RTShader;

/************************************************************************/
/*                                                                      */
/************************************************************************/
String TerrainTransform::Type = "TerrainTransform";

bool TerrainTransform::preAddToRenderState(const RenderState* renderState, Pass* srcPass, Pass* dstPass)
{
    auto terrainAny = srcPass->getUserObjectBindings().getUserAny("Terrain");
    mTerrain = any_cast<const Terrain*>(terrainAny);
    mCompressed = mTerrain->_getUseVertexCompression();
    mAlign = mTerrain->getAlignment();
    return true;
}

void TerrainTransform::updateParams()
{
    if(!mCompressed)
        return;

    mPointTrans->setGpuParameter(mTerrain->getPointTransform());

    float baseUVScale = 1.0f / (mTerrain->getSize() - 1);
    mBaseUVScale->setGpuParameter(baseUVScale);
}

//-----------------------------------------------------------------------
bool TerrainTransform::createCpuSubPrograms(ProgramSet* programSet)
{
    static Operand::OpMask heightAxis[3] = {Operand::OPM_Y, Operand::OPM_Z, Operand::OPM_X};

    Program* vsProgram = programSet->getCpuProgram(GPT_VERTEX_PROGRAM);
    Function* vsEntry = vsProgram->getEntryPointFunction();

    auto posType = mCompressed ? GCT_INT2 : GCT_FLOAT4;

    auto wvpMatrix = vsProgram->resolveParameter(GpuProgramParameters::ACT_WORLDVIEWPROJ_MATRIX);

    auto lodMorph = vsProgram->resolveAutoParameterInt(GpuProgramParameters::ACT_CUSTOM, GCT_FLOAT2, Terrain::LOD_MORPH_CUSTOM_PARAM);

    auto positionIn = vsEntry->resolveInputParameter(Parameter::SPC_POSITION_OBJECT_SPACE, posType);
    auto positionOut = vsEntry->resolveOutputParameter(Parameter::SPC_POSITION_PROJECTIVE_SPACE);
    auto uv = vsEntry->resolveOutputParameter(Parameter::SPC_TEXTURE_COORDINATE0, GCT_FLOAT2);
    auto delta = vsEntry->resolveInputParameter(Parameter::SPC_TEXTURE_COORDINATE1, GCT_FLOAT2);

    // Add dependency.
    vsProgram->addDependency("FFPLib_Transform");
    vsProgram->addDependency("TerrainTransforms");

    auto stage = vsEntry->getStage(FFP_VS_TRANSFORM);

    if(mCompressed)
    {
        mPointTrans = vsProgram->resolveParameter(GCT_MATRIX_4X4, "pointTrans");
        mBaseUVScale = vsProgram->resolveParameter(GCT_FLOAT1, "baseUVScale");
        auto height = vsEntry->resolveInputParameter(Parameter::SPC_TEXTURE_COORDINATE0, GCT_FLOAT1);
        auto position = vsEntry->resolveLocalParameter(Parameter::SPC_POSITION_OBJECT_SPACE);

        stage.callFunction("expandVertex",
                       {In(mPointTrans), In(mBaseUVScale), In(positionIn), In(height), Out(position), Out(uv)});
        positionIn = position;
    }
    else
    {
        auto uvin = vsEntry->resolveInputParameter(Parameter::SPC_TEXTURE_COORDINATE0, GCT_FLOAT2);
        stage.assign(In(uvin).xy(), uv);
    }

    stage.callBuiltin("mul", wvpMatrix, positionIn, positionOut);
    stage.callFunction("applyLODMorph", {In(delta), In(lodMorph), InOut(positionOut).mask(heightAxis[mAlign])});

    return true;
}

String TerrainSurface::Type = "TerrainSurface";

bool TerrainSurface::setParameter(const String& name, const String& value)
{
    if(name == "for_composite_map")
    {
        return StringConverter::parse(value, mForCompositeMap);
    }
    else if (name == "use_parallax_mapping")
    {
        return StringConverter::parse(value, mUseParallaxMapping);
    }
    else if (name == "use_parallax_occlusion_mapping")
    {
        return StringConverter::parse(value, mUseParallaxOcclusionMapping);
    }
    else if (name == "use_specular_mapping")
    {
        return StringConverter::parse(value, mUseSpecularMapping);
    }
    else if (name == "use_normal_mapping")
    {
        return StringConverter::parse(value, mUseNormalMapping);
    }

    return false;
}

bool TerrainSurface::preAddToRenderState(const RenderState* renderState, Pass* srcPass, Pass* dstPass)
{
    mTerrain = any_cast<const Terrain*>(srcPass->getUserObjectBindings().getUserAny("Terrain"));

    SamplerPtr clampSampler = TextureManager::getSingleton().createSampler();
    clampSampler->setAddressingMode(TAM_CLAMP);
    clampSampler->setFiltering(FT_MIP, FO_NONE);

    auto tu = dstPass->createTextureUnitState();
    tu->setTexture(mTerrain->getTerrainNormalMap());
    tu->setSampler(clampSampler);

    if (auto cm = mTerrain->getGlobalColourMap())
    {
        tu = dstPass->createTextureUnitState();
        tu->setTexture(cm);
        tu->setSampler(clampSampler);
    }

    if(auto lm = mTerrain->getLightmap())
    {
        tu = dstPass->createTextureUnitState();
        tu->setTexture(lm);
        tu->setSampler(clampSampler);
    }

    for(auto bt : mTerrain->getBlendTextures())
    {
        tu = srcPass->createTextureUnitState();
        tu->setTexture(bt);
        tu->setSampler(clampSampler);
    }

    mUVMul.resize((mTerrain->getLayerCount() + 3) / 4); // integer ceil

    mUseNormalMapping = mUseNormalMapping && !mTerrain->getLayerTextureName(0, 1).empty();
    for (int i = 0; i < mTerrain->getLayerCount(); ++i)
    {
        srcPass->createTextureUnitState(mTerrain->getLayerTextureName(i, 0));
        if (mUseNormalMapping)
            srcPass->createTextureUnitState(mTerrain->getLayerTextureName(i, 1));
    }

    return true;
}

void TerrainSurface::updateParams()
{
    for (size_t i = 0; i < mUVMul.size(); i++)
    {
        Vector4 uvMul(mTerrain->getLayerUVMultiplier(i + 0), mTerrain->getLayerUVMultiplier(i + 1),
                      mTerrain->getLayerUVMultiplier(i + 2), mTerrain->getLayerUVMultiplier(i + 3));
        mUVMul[i]->setGpuParameter(uvMul);
    }
}

static Operand::OpMask channel[4] = {Operand::OPM_X, Operand::OPM_Y, Operand::OPM_Z, Operand::OPM_W};

//-----------------------------------------------------------------------
bool TerrainSurface::createCpuSubPrograms(ProgramSet* programSet)
{
    Program* vsProgram = programSet->getCpuProgram(GPT_VERTEX_PROGRAM);
    Program* psProgram = programSet->getCpuProgram(GPT_FRAGMENT_PROGRAM);
    Function* psMain = psProgram->getMain();
    Function* vsMain = vsProgram->getMain();

    psProgram->addDependency("FFPLib_Transform");
    psProgram->addDependency("SGXLib_NormalMap");
    psProgram->addDependency("SGXLib_IntegratedPSSM");
    psProgram->addDependency("TerrainSurface");

    if(mUseNormalMapping)
        psProgram->addPreprocessorDefines("TERRAIN_NORMAL_MAPPING");

    auto uvVS = vsMain->resolveOutputParameter(Parameter::SPC_TEXTURE_COORDINATE0, GCT_FLOAT2);

    if (mForCompositeMap)
    {
        auto uvIn = vsMain->resolveInputParameter(Parameter::SPC_TEXTURE_COORDINATE0, GCT_FLOAT2);
        // forward texcoords in VS
        vsMain->getStage(FFP_VS_TEXTURING).assign(uvIn, uvVS);
    }

    ParameterPtr viewPos;
    if (mUseNormalMapping && mUseParallaxMapping)
    {
        psProgram->addPreprocessorDefines("TERRAIN_PARALLAX_MAPPING");
        if (mUseParallaxOcclusionMapping)
        {
            psProgram->addPreprocessorDefines("POM_MAX_DISTANCE=400.0,POM_LAYER_COUNT=32");
        }

        // assuming: lighting stage computed this
        auto vsOutViewPos = vsMain->resolveOutputParameter(Parameter::SPC_POSITION_VIEW_SPACE);
        viewPos = psMain->resolveInputParameter(vsOutViewPos);
    }

    auto uvPS = psMain->resolveInputParameter(uvVS);
    uvPS->setHighP(true);

    for(auto& uvMul : mUVMul)
    {
        uvMul = psProgram->resolveParameter(GCT_FLOAT4, "uvMul");
    }

    int texUnit = 0;
    auto globalNormal = psProgram->resolveParameter(GCT_SAMPLER2D, "globalNormal", texUnit++);

    ParameterPtr globalColourMap;
    if (mTerrain->getGlobalColourMap())
        globalColourMap = psProgram->resolveParameter(GCT_SAMPLER2D, "globalColour", texUnit++);

    ParameterPtr lightMap;
    if(mTerrain->getLightmap())
        lightMap = psProgram->resolveParameter(GCT_SAMPLER2D, "lightMap", texUnit++);

    auto normal = psMain->resolveLocalParameter(Parameter::SPC_NORMAL_VIEW_SPACE);
    auto ITMat = psProgram->resolveParameter(GpuProgramParameters::ACT_NORMAL_MATRIX);

    auto diffuse = psMain->resolveLocalParameter(Parameter::SPC_COLOR_DIFFUSE);
    auto diffuseSpec = psMain->resolveLocalParameter(GCT_FLOAT4, "diffuseSpec");
    auto TSnormal = psMain->resolveLocalParameter(GCT_FLOAT3, "TSnormal");
    auto texTmp = psMain->resolveLocalParameter(GCT_FLOAT4, "texTmp");

    auto outDiffuse = psMain->resolveOutputParameter(Parameter::SPC_COLOR_DIFFUSE);

    auto stage = psMain->getStage(FFP_PS_COLOUR_BEGIN);
    stage.assign(Vector4(1), outDiffuse); // FFPColour logic
    stage.callFunction("SGX_FetchNormal", globalNormal, uvPS, normal);
    stage.callBuiltin("mul", ITMat, normal, normal);

    auto psSpecular = psMain->resolveLocalParameter(Parameter::SPC_COLOR_SPECULAR);
    stage.assign(Vector4::ZERO, psSpecular);

    std::vector<ParameterPtr> blendWeights;
    for(auto bt : mTerrain->getBlendTextures())
    {
        auto weight = psMain->resolveLocalParameter(GCT_FLOAT4, StringUtil::format("blendWeight%d", texUnit));
        auto blendTex = psProgram->resolveParameter(GCT_SAMPLER2D, "blendTex", texUnit++);
        stage.sampleTexture(blendTex, uvPS, weight);
        blendWeights.push_back(weight);
    }

    // Call TBN calculation
    auto psOutTBN = psMain->resolveLocalParameter(GpuConstantType::GCT_MATRIX_3X3, "TBN");
    stage.callFunction("SGX_CalculateTerrainTBN", {In(normal), In(ITMat), Out(psOutTBN)});

    stage.assign(Vector4::ZERO, diffuseSpec);
    stage.assign(Vector3(0, 0, 1), TSnormal);
    for (int l = 0; l < mTerrain->getLayerCount(); ++l)
    {
        auto blendWeight = l == 0 ? In(1.0f) : In(blendWeights[(l - 1) / 4]).mask(channel[(l - 1) % 4]);
        auto difftex = psProgram->resolveParameter(GCT_SAMPLER2D, "difftex", texUnit++);
        std::vector<Operand> args = {blendWeight, In(uvPS), In(mUVMul[l/4]).mask(channel[l % 4])};
        if (mUseNormalMapping)
        {
            if (mUseParallaxMapping)
            {
                args.push_back(In(viewPos));
                args.push_back(In(0.04)); //Scale
                args.push_back(In(psOutTBN));
            }

            auto normtex = psProgram->resolveParameter(GCT_SAMPLER2D, "normtex", texUnit++);
            args.push_back(In(normtex));
            args.push_back(Out(TSnormal));
        }
        args.push_back(In(difftex));
        args.push_back(Out(diffuseSpec));
        stage.callFunction("blendTerrainLayer", {args});
    }

    if(mUseNormalMapping)
        stage.callFunction("transformToTS", {In(TSnormal), In(ITMat), InOut(normal)});

    // fake vertexcolour input for TVC_SPECULAR
    if(mUseSpecularMapping)
        stage.mul(In(diffuseSpec).w(), Vector4(1), diffuse);

    if(lightMap)
    {
        auto shadowFactor = psMain->getLocalParameter("lShadowFactor");
        if(!shadowFactor)
        {
            shadowFactor = psMain->resolveLocalParameter(GCT_FLOAT1, "lShadowFactor", 1);
            psProgram->addPreprocessorDefines("SHADOWLIGHT_COUNT=1");
        }
        stage = psMain->getStage(FFP_PS_COLOUR_BEGIN - 1); // before the PSSM stage
        stage.assign({In(1), Out(shadowFactor), At(0)});

        stage = psMain->getStage(FFP_PS_COLOUR_BEGIN + 1); // after the PSSM stage
        stage.callFunction("getShadowFactor", {In(lightMap), In(uvPS), InOut(shadowFactor), At(0)});
    }

    if(globalColourMap)
    {
        stage = psMain->getStage(FFP_PS_COLOUR_BEGIN + 2); // after lighting calculations
        stage.sampleTexture(globalColourMap, uvPS, texTmp);
        stage.mul(In(diffuseSpec).xyz(), In(texTmp).xyz(), Out(diffuseSpec).xyz());
    }

    stage = psMain->getStage(FFP_PS_TEXTURING);
    stage.mul(diffuseSpec, outDiffuse, outDiffuse);

    // FFPColour logic
    psMain->getStage(FFP_PS_COLOUR_END)
        .add(In(outDiffuse).xyz(), In(psSpecular).xyz(), Out(outDiffuse).xyz());
    return true;
}

//-----------------------------------------------------------------------
const String& TerrainTransformFactory::getType() const
{
    return TerrainTransform::Type;
}
SubRenderState* TerrainTransformFactory::createInstanceImpl()
{
    return OGRE_NEW TerrainTransform();
}

const String& TerrainSurfaceFactory::getType() const
{
    return TerrainSurface::Type;
}
SubRenderState* TerrainSurfaceFactory::createInstanceImpl()
{
    return OGRE_NEW TerrainSurface();
}

}