/* osgEarth
 * Copyright 2025 Pelican Mapping
 * MIT License
 */
#ifndef OSGEARTH_SHADER_UTILS_H
#define OSGEARTH_SHADER_UTILS_H 1

#include <osgEarth/Common>
#include <osg/NodeCallback>
#include <osg/StateSet>
#include <osg/Uniform>
#include <osg/Light>
#include <osg/Material>
#include <osg/observer_ptr>
#include <osgDB/Options>
#include <map>
#include <unordered_map>
#include <functional>

namespace osgEarth
{
    /**
     * ShaderPolicy encodes general behavior when deciding how to
     * employ shaders in certain situations
     */
    enum ShaderPolicy
    {
        SHADERPOLICY_DISABLE,
        SHADERPOLICY_GENERATE,
        SHADERPOLICY_INHERIT
    };
}

namespace osgEarth { namespace Util
{
    struct OSGEARTH_EXPORT ShaderUtils
    {
        // Installs a default, one-texture shader on a stateset
        // for basic rendering.
        static void installDefaultShader(osg::StateSet*);
    };

    /**
     * Preprocesses GLES shader source to include our osg_LightProducts and osg_LightSourceParameters
     * definitions and uniforms.
     */
    class OSGEARTH_EXPORT ShaderPreProcessor
    {
    public:

        static void runPre(std::string& source);

        static void runPost(osg::Shader* shader);

        struct PreCallbackInfo {
            osg::observer_ptr<osg::Referenced> host;
            std::function<void(std::string& source, osg::Referenced* host)> function;
        };

        struct PostCallbackInfo {
            osg::observer_ptr<osg::Referenced> host;
            std::function<void(osg::Shader*, osg::Referenced* host)> function;
        };

        static std::unordered_map<
            UID,
            PreCallbackInfo> _pre_callbacks;

        static std::unordered_map<
            UID,
            PostCallbackInfo> _post_callbacks;
    };

    /**
     * Helper class for dealing with array uniforms. Array uniform naming works
     * differently on different drivers (ATI vs NVIDIA), so this class helps mitigate
     * those differences.
     */
    class OSGEARTH_EXPORT ArrayUniform //  : public osg::Referenced
    {
    public:
        /** Empty array uniform */
        ArrayUniform() { }

        /**
         * Creates or retrieves a named uniform array.
         */
        ArrayUniform(
            const std::string& name,
            osg::Uniform::Type type,
            osg::StateSet*     stateSet,
            unsigned           size =1 );

        /** dtor */
        virtual ~ArrayUniform() { }

        void attach(
            const std::string& name,
            osg::Uniform::Type type,
            osg::StateSet*     stateSet,
            unsigned           size =1 );

        void detach();

        void setElement( unsigned index, int value );
        void setElement( unsigned index, unsigned value );
        void setElement( unsigned index, bool value );
        void setElement( unsigned index, float value );
        void setElement( unsigned index, const osg::Matrixf& value );
        void setElement( unsigned index, const osg::Vec3f& value );
        void setElement( unsigned index, const osg::Vec4f& value);

        bool getElement( unsigned index, int& out_value ) const;
        bool getElement( unsigned index, unsigned& out_value ) const;
        bool getElement( unsigned index, bool& out_value ) const;
        bool getElement( unsigned index, float& out_value ) const;
        bool getElement( unsigned index, osg::Matrixf& out_value ) const;
        bool getElement( unsigned index, osg::Vec3f& out_value ) const;
        bool getElement( unsigned index, osg::Vec4f& out_value) const;

        bool isValid() const { return _uniform.valid() && _uniformAlt.valid(); }
        int getNumElements() const { return isValid() ? _uniform->getNumElements() : -1; }

        bool isDirty() const { return
            (_uniform.valid() && _uniform->getModifiedCount() > 0) ||
            (_uniformAlt.valid() && _uniformAlt->getModifiedCount() > 0); }

    private:
        osg::ref_ptr<osg::Uniform>       _uniform;
        osg::ref_ptr<osg::Uniform>       _uniformAlt;
        osg::observer_ptr<osg::StateSet> _stateSet;

        void ensureCapacity( unsigned newSize );
    };

    /**
     * Shader that will discard fragments whose alpha component falls below
     * the specified threshold.
     */
    class OSGEARTH_EXPORT DiscardAlphaFragments
    {
    public:
        void install(osg::StateSet* ss, float minAlpha) const;
        void uninstall(osg::StateSet* ss) const;
    };

    /**
     * Class to parse messages from an info log.
     */
    class OSGEARTH_EXPORT ShaderInfoLog
    {
    public:
        ShaderInfoLog(
            const osg::Program* program,
            const std::string& log);

        void dumpErrors(
            osg::State&) const;

    private:
        osg::ref_ptr<const osg::Program> _program;
        std::string _log;

        struct Message {
            unsigned line;
            unsigned column;
            std::string message;
        };
    };
} }

#endif // OSGEARTH_SHADER_UTILS_H
