#include "material.hpp"

#include <array>
#include <osg/BlendFunc>
#include <osg/Capability>
#include <osg/Depth>
#include <osg/Fog>
#include <osg/Texture2D>

#include <components/resource/scenemanager.hpp>
#include <components/sceneutil/depth.hpp>
#include <components/sceneutil/texmat.hpp>
#include <components/sceneutil/util.hpp>
#include <components/shader/shadermanager.hpp>
#include <components/stereo/stereomanager.hpp>

#include <mutex>

namespace
{
    class BlendmapTexMat
    {
    public:
        static const osg::ref_ptr<osg::Uniform>& value(const int blendmapScale)
        {
            static BlendmapTexMat instance;
            return instance.get(static_cast<float>(blendmapScale));
        }

        const osg::ref_ptr<osg::Uniform>& get(const float blendmapScale)
        {
            const std::lock_guard<std::mutex> lock(mMutex);
            auto texMat = mTexMatMap.find(blendmapScale);
            if (texMat == mTexMatMap.end())
            {
                osg::Matrixf matrix;
                float scale = (blendmapScale / (blendmapScale + 1.f));
                matrix.preMultTranslate(osg::Vec3f(0.5f, 0.5f, 0.f));
                matrix.preMultScale(osg::Vec3f(scale, scale, 1.f));
                matrix.preMultTranslate(osg::Vec3f(-0.5f, -0.5f, 0.f));
                // We need to nudge the blendmap to look like vanilla.
                // This causes visible seams unless the blendmap's resolution is doubled, but Vanilla also doubles the
                // blendmap, apparently.
                matrix.preMultTranslate(osg::Vec3f(1.0f / blendmapScale / 4.0f, -1.0f / blendmapScale / 4.0f, 0.f));

                texMat = mTexMatMap.emplace(blendmapScale, SceneUtil::createTexMatUniform(1, matrix)).first;
            }
            return texMat->second;
        }

    private:
        std::mutex mMutex;
        std::map<float, osg::ref_ptr<osg::Uniform>> mTexMatMap;
    };

    class LayerTexMat
    {
    public:
        static const osg::ref_ptr<osg::Uniform>& value(const float layerTileSize)
        {
            static LayerTexMat instance;
            return instance.get(layerTileSize);
        }

        const osg::ref_ptr<osg::Uniform>& get(const float layerTileSize)
        {
            const std::lock_guard<std::mutex> lock(mMutex);
            auto texMat = mTexMatMap.find(layerTileSize);
            if (texMat == mTexMatMap.end())
            {
                texMat = mTexMatMap
                             .insert(std::make_pair(layerTileSize,
                                 SceneUtil::createTexMatUniform(
                                     0, osg::Matrix::scale(osg::Vec3f(layerTileSize, layerTileSize, 1.f)))))
                             .first;
            }
            return texMat->second;
        }

    private:
        std::mutex mMutex;
        std::map<float, osg::ref_ptr<osg::Uniform>> mTexMatMap;
    };

    class EqualDepth
    {
    public:
        static const osg::ref_ptr<osg::Depth>& value()
        {
            static EqualDepth instance;
            return instance.mValue;
        }

    private:
        osg::ref_ptr<osg::Depth> mValue;

        EqualDepth()
            : mValue(new SceneUtil::AutoDepth)
        {
            mValue->setFunction(osg::Depth::EQUAL);
        }
    };

    class LequalDepth
    {
    public:
        static const osg::ref_ptr<osg::Depth>& value()
        {
            static LequalDepth instance;
            return instance.mValue;
        }

    private:
        osg::ref_ptr<osg::Depth> mValue;

        LequalDepth()
            : mValue(new SceneUtil::AutoDepth(osg::Depth::LEQUAL))
        {
        }
    };

    class BlendFuncFirst
    {
    public:
        static const osg::ref_ptr<osg::BlendFunc>& value()
        {
            static BlendFuncFirst instance;
            return instance.mValue;
        }

    private:
        osg::ref_ptr<osg::BlendFunc> mValue;

        BlendFuncFirst()
            : mValue(new osg::BlendFunc(osg::BlendFunc::SRC_ALPHA, osg::BlendFunc::ZERO))
        {
        }
    };

    class BlendFunc
    {
    public:
        static const osg::ref_ptr<osg::BlendFunc>& value()
        {
            static BlendFunc instance;
            return instance.mValue;
        }

    private:
        osg::ref_ptr<osg::BlendFunc> mValue;

        BlendFunc()
            : mValue(new osg::BlendFunc(osg::BlendFunc::SRC_ALPHA, osg::BlendFunc::ONE))
        {
        }
    };

    class UniformCollection
    {
    public:
        static const UniformCollection& value()
        {
            static UniformCollection instance;
            return instance;
        }

        osg::ref_ptr<osg::Uniform> mDiffuseMap;
        osg::ref_ptr<osg::Uniform> mBlendMap;
        osg::ref_ptr<osg::Uniform> mNormalMap;
        osg::ref_ptr<osg::Uniform> mColorMode;

        UniformCollection()
            : mDiffuseMap(new osg::Uniform("diffuseMap", 0))
            , mBlendMap(new osg::Uniform("blendMap", 1))
            , mNormalMap(new osg::Uniform("normalMap", 2))
            , mColorMode(new osg::Uniform("colorMode", 2))
        {
        }
    };

    const std::array<Shader::ShaderManager::SamplerBindingMap, 5>& getSinglePassSamplers()
    {
        static const auto sSinglePassSamplers = []() {
            std::array<Shader::ShaderManager::SamplerBindingMap, 5> maps;
            for (int numLayers = 2; numLayers <= 6; ++numLayers)
            {
                auto& m = maps[numLayers - 2];
                for (int i = 0; i < numLayers; ++i)
                {
                    m["diffuseMap" + std::to_string(i)] = i;
                    m["blendMap" + std::to_string(i)] = numLayers + i;
                }
            }
            return maps;
        }();
        return sSinglePassSamplers;
    }
}

namespace Terrain
{
    std::vector<osg::ref_ptr<osg::StateSet>> createPasses(Resource::SceneManager* sceneManager,
        const std::vector<TextureLayer>& layers, const std::vector<osg::ref_ptr<osg::Texture2D>>& blendmaps,
        int blendmapScale, float layerTileSize, bool isComposite, bool esm4terrain)
    {
        auto& shaderManager = sceneManager->getShaderManager();
        std::vector<osg::ref_ptr<osg::StateSet>> passes;

        bool canUseSinglePass = !isComposite && !blendmaps.empty() && layers.size() >= 2 && layers.size() <= 6
            && (static_cast<int>(layers.size() * 2) <= shaderManager.getAvailableTextureUnits())
            && (blendmaps.size() == layers.size());

        if (canUseSinglePass)
        {
            for (const auto& layer : layers)
            {
                if (!layer.mDiffuseMap || layer.mNormalMap || layer.mSpecular || layer.mParallax)
                {
                    canUseSinglePass = false;
                    break;
                }
            }
            for (const auto& bm : blendmaps)
            {
                if (!bm)
                {
                    canUseSinglePass = false;
                    break;
                }
            }
        }

        if (canUseSinglePass)
        {
            osg::ref_ptr<osg::StateSet> stateset(new osg::StateSet);

            stateset->setRenderBinDetails(0, "RenderBin");
            stateset->setAttributeAndModes(LequalDepth::value(), osg::StateAttribute::ON);
            stateset->setMode(GL_BLEND, osg::StateAttribute::OFF);

            if (layerTileSize != 1.f)
                stateset->addUniform(LayerTexMat::value(layerTileSize));

            if (!esm4terrain)
                stateset->addUniform(BlendmapTexMat::value(blendmapScale));

            stateset->addUniform(UniformCollection::value().mColorMode);

            const size_t numLayers = layers.size();
            for (size_t i = 0; i < numLayers; ++i)
            {
                stateset->setTextureAttribute(static_cast<int>(i), layers[i].mDiffuseMap);

                int blendUnit = static_cast<int>(numLayers + i);
                stateset->setTextureAttribute(blendUnit, blendmaps[i].get());
            }

            Shader::ShaderManager::DefineMap defineMap;
            defineMap["singlePassTerrain"] = "1";
            defineMap["numLayers"] = std::to_string(numLayers);
            defineMap["normalMap"] = "0";
            defineMap["blendMap"] = "0";
            defineMap["specularMap"] = "0";
            defineMap["parallax"] = "0";
            defineMap["reconstructNormalZ"] = "0";
            Stereo::shaderStereoDefines(defineMap);

            const auto& samplers = getSinglePassSamplers()[numLayers - 2];
            stateset->setAttributeAndModes(shaderManager.getProgram("terrain", defineMap, nullptr, samplers));

            passes.push_back(stateset);
            return passes;
        }

        unsigned int blendmapIndex = 0;
        for (std::vector<TextureLayer>::const_iterator it = layers.begin(); it != layers.end(); ++it)
        {
            bool firstLayer = (it == layers.begin());

            osg::ref_ptr<osg::StateSet> stateset(new osg::StateSet);

            if (!blendmaps.empty())
            {
                stateset->setMode(GL_BLEND, osg::StateAttribute::ON);
                // This pass needs normals-buffer blending.
                stateset->setAttribute(new osg::Enablei(GL_BLEND, 1));
                stateset->setRenderBinDetails(firstLayer ? 0 : 1, "RenderBin");
                if (!firstLayer)
                {
                    stateset->setAttributeAndModes(BlendFunc::value(), osg::StateAttribute::ON);
                    stateset->setAttributeAndModes(EqualDepth::value(), osg::StateAttribute::ON);
                }
                else
                {
                    stateset->setAttributeAndModes(BlendFuncFirst::value(), osg::StateAttribute::ON);
                    stateset->setAttributeAndModes(LequalDepth::value(), osg::StateAttribute::ON);
                }
            }

            stateset->setTextureAttribute(0, it->mDiffuseMap);
            stateset->addUniform(UniformCollection::value().mDiffuseMap);

            if (layerTileSize != 1.f)
                stateset->addUniform(LayerTexMat::value(layerTileSize));

            if (!blendmaps.empty())
            {
                osg::ref_ptr<osg::Texture2D> blendmap = blendmaps.at(blendmapIndex++);

                stateset->setTextureAttribute(1, blendmap.get());
                if (!esm4terrain)
                    stateset->addUniform(BlendmapTexMat::value(blendmapScale));
                stateset->addUniform(UniformCollection::value().mBlendMap);
            }
            if (isComposite)
            {
                stateset->setAttributeAndModes(
                    shaderManager.getProgram("terrain_composite", { { "blendMap", !blendmaps.empty() ? "1" : "0" } }));
            }
            else
            {
                bool parallax = it->mNormalMap && it->mParallax;
                bool reconstructNormalZ = false;

                if (it->mNormalMap)
                {
                    stateset->setTextureAttribute(2, it->mNormalMap);
                    stateset->addUniform(UniformCollection::value().mNormalMap);

                    // Special handling for red-green normal maps (e.g. BC5 or R8G8).
                    const osg::Image* image = it->mNormalMap->getImage(0);
                    if (image)
                    {
                        switch (SceneUtil::computeUnsizedPixelFormat(image->getPixelFormat()))
                        {
                            case GL_RG:
                            case GL_RG_INTEGER:
                            {
                                reconstructNormalZ = true;
                                parallax = false;
                            }
                        }
                    }
                }

                Shader::ShaderManager::DefineMap defineMap;
                defineMap["normalMap"] = (it->mNormalMap) ? "1" : "0";
                defineMap["blendMap"] = (!blendmaps.empty()) ? "1" : "0";
                defineMap["specularMap"] = it->mSpecular ? "1" : "0";
                defineMap["parallax"] = parallax ? "1" : "0";
                defineMap["reconstructNormalZ"] = reconstructNormalZ ? "1" : "0";
                Stereo::shaderStereoDefines(defineMap);

                stateset->setAttributeAndModes(shaderManager.getProgram("terrain", defineMap));
                stateset->addUniform(UniformCollection::value().mColorMode);
            }

            passes.push_back(stateset);
        }
        return passes;
    }

}