/* -*-c++-*- */
/* osgEarth - Geospatial SDK for OpenSceneGraph
* Copyright 2018 Pelican Mapping
* http://osgearth.org
*
* osgEarth is free software; you can redistribute it and/or modify
* it under the terms of the GNU Lesser General Public License as published by
* the Free Software Foundation; either version 2 of the License, or
* (at your option) any later version.
*
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
* FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
* IN THE SOFTWARE.
*
* You should have received a copy of the GNU Lesser General Public License
* along with this program.  If not, see <http://www.gnu.org/licenses/>
*/
#ifndef OSGEARTH_VIRTUAL_PROGRAM_H
#define OSGEARTH_VIRTUAL_PROGRAM_H 1

#include <osgEarth/Common>
#include <osgEarth/ThreadingUtils>
#include <osgEarth/Containers>
#include <osgEarth/optional>
#include <osg/Shader>
#include <osg/Program>
#include <osg/StateAttribute>
#include <osg/buffered_value>
#include <string>
#include <map>

#if defined(OSG_GLES2_AVAILABLE)
#    define GLSL_VERSION                 100
#    define GLSL_VERSION_STR             "100"
#    define GLSL_DEFAULT_PRECISION_FLOAT "precision highp float;\nprecision highp int;"
#elif defined(OSG_GLES3_AVAILABLE)
#    define GLSL_VERSION                 300
#    define GLSL_VERSION_STR             "300 es"
#    define GLSL_DEFAULT_PRECISION_FLOAT "precision highp float;\nprecision highp int;"
#elif defined(OSG_GL3_AVAILABLE)
#    define GLSL_VERSION                 330
#    define GLSL_VERSION_STR             "330"
#    define GLSL_DEFAULT_PRECISION_FLOAT ""
#else
#    define GLSL_VERSION                 330
#    define GLSL_VERSION_STR             "330 compatibility" 
#    define GLSL_DEFAULT_PRECISION_FLOAT ""
#endif

namespace osgEarth
{
    namespace ShaderComp
    {
        // User function injection points.
        enum FunctionLocation
        {
            // vertex is in model space (equivalent to gl_Vertex).
            LOCATION_VERTEX_MODEL,

            // vertex is in view(aka eye) coordinates, with the camera at 0,0,0 
            // looking down the -Z axis.
            LOCATION_VERTEX_VIEW,

            // vertex is in post-perspective coordinates; [-w..w] along each axis
            LOCATION_VERTEX_CLIP,

            // tessellation control shader; model space
            LOCATION_TESS_CONTROL,

            // tessellation evalulation shader; model space
            LOCATION_TESS_EVALUATION,

            // geometry shader; inputs are in model space.
            LOCATION_GEOMETRY,

            // fragment is being colored.
            LOCATION_FRAGMENT_COLORING,

            // fragment is being lit.
            LOCATION_FRAGMENT_LIGHTING,

            // fragment output is being assigned.
            LOCATION_FRAGMENT_OUTPUT,

            // not defined.
            LOCATION_UNDEFINED
        };

        /**
         * Callback that accepts a user-injected shader function (set with
         * setFunction) for inclusing in the program at render time.
         * @deprecated (remove after user support)
         */
        class AcceptCallback : public osg::Referenced
        {
        public:
            // implement this to accept or reject based on the state
            virtual bool operator()(const osg::State& state) =0;
        protected:
            virtual ~AcceptCallback() { }
        };

        // User function (used internally)
        struct Function
        {
            std::string                  _name;
            osg::ref_ptr<AcceptCallback> _accept; // @deprecated (remove after user support)
            
            // @dperecated (remove after user support)
            bool accept(const osg::State& state) const {
                return (!_accept.valid()) || (_accept->operator()(state) == true);
            }
        };

        typedef std::pair<float, Function> OrderedFunction;

        // ordered set of user functions.
        typedef std::multimap<float, Function> OrderedFunctionMap; // duplicate keys allowed

        // user function sets, categorized by function location.
        typedef std::map<FunctionLocation, OrderedFunctionMap> FunctionLocationMap;

        // Mask values that represents which stages a composed shader contains.
        enum StageMaskValues
        {
            STAGE_VERTEX          = 1 << 0,
            STAGE_TESSCONTROL     = 1 << 1,
            STAGE_TESSEVALUATION  = 1 << 2,
            STAGE_GEOMETRY        = 1 << 3,
            STAGE_FRAGMENT        = 1 << 4,
            STAGE_COMPUTE         = 1 << 5
        };
        typedef unsigned StageMask;
    }


    /**
     * A Shader that can compile into different stages of the program pipeline
     * depending on which stages are in use. For example, you may designate a
     * VP component to run in the "VERTEX_CLIP" location. But if a geometry shader
     * is present, this function must be moved to the end of the GEOMETRY stage
     * instead. PolyShader supports this. It also preserves the original shader
     * source in case someone wants to access it via VirtualProgram::getShaders
     * (for example, to call VP components from an external shader program).
     */
    class OSGEARTH_EXPORT PolyShader : public osg::Referenced
    {
    public:
        /** Construct a polyshader */
        PolyShader();

        /** Construct a polyshader */
        PolyShader(osg::Shader* shader);

        /** Name of the polyshader */
        void setName(const std::string& name) { _name = name; }
        const std::string& getName() const { return _name; }

        /** Nominal stage at which this shader would run */
        void setLocation(ShaderComp::FunctionLocation loc);
        ShaderComp::FunctionLocation getLocation() const { return _location; }

        /** Shader source code (unprocessed) */
        void setShaderSource(const std::string& source);
        const std::string& getShaderSource() const { return _source; }

        /** Given a mask of all available stages, return the shader that will run
            this polyshader in the appropriate stage. */
        osg::Shader* getShader(ShaderComp::StageMask mask) const;

        /** The shader in its named stage. */
        osg::Shader* getNominalShader() const { return _nominalShader.get(); }

        /** Generates the shaders. */
        void prepare();

        /** Called from the draw context to resize shader buffers as necessary (OSG) */
        virtual void resizeGLObjectBuffers(unsigned maxSize);

        virtual void releaseGLObjects(osg::State* state) const;

        static PolyShader * lookUpShader(const std::string & functionName, const std::string & shaderSource, ShaderComp::FunctionLocation location);

    protected:
        virtual ~PolyShader() { }

        std::string                  _name;
        std::string                  _source;
        ShaderComp::FunctionLocation _location;

        // Originally specified stage for this shader
        osg::Shader::Type            _nominalType;

        // The shader built for the nominal program stage (nominalType)
        osg::ref_ptr<osg::Shader>    _nominalShader;

        // Same shader, but set up for injection into a different stage
        osg::ref_ptr<osg::Shader>    _geomShader;
        osg::ref_ptr<osg::Shader>    _tessevalShader;

        bool                         _dirty;

        static Threading::Mutex _cacheMutex;
        typedef std::map<std::pair<std::string, std::string>, osg::ref_ptr<PolyShader> > PolyShaderCache;
        static PolyShaderCache _polyShaderCache;

    };

    //typedef std::vector< osg::ref_ptr<PolyShader> > ProgramKey;
    typedef std::vector<PolyShader*> ProgramKey;


    class /*internal*/ ProgramRepo : public osg::Referenced
    {
    public:

        // Referenced, b/c more than one ProgramKey can point to the same Entry
        // (shared program)
        struct Entry : public osg::Referenced
        {
            osg::ref_ptr<osg::Program> _program;
            unsigned                   _frameLastUsed;
            std::set<UID>              _users;
        };

        typedef std::map<ProgramKey, osg::ref_ptr<Entry> > ProgramMap;

        //! Exclusive lock on the repo - enclose any calls to this object
        //! with a lock/unlock pair
        void lock();

        //! unlock the repo
        void unlock();

        //! Search for a program matching the key and return it, adding the user
        //! to its users list and updating the frame number.
        osg::ref_ptr<osg::Program> use(const ProgramKey& key, unsigned frameNumber, UID user);

        //! Insert a new program into the repo
        void add(const ProgramKey& key, osg::ref_ptr<osg::Program>& inOut, unsigned frameNumber, UID user);

        //! Release anything used by this user
        void release(UID user, osg::State* state);

        //! Prune expired data
        void prune(unsigned frameNumber, osg::State* state);

        void resizeGLObjectBuffers(unsigned maxSize);

        void releaseGLObjects(osg::State* state) const;

        ~ProgramRepo();

    private:
        mutable ProgramMap _db;
        mutable Threading::Mutex _m;
    };


    /**
     * VirtualProgram enables GLSL shader composition within osgEarth. It automatically
     * assembles shader functions into a full shader program at run time. You can add
     * or remove functions (injection points) at any time.
     *
     * Read about shader composition:
     * http://docs.osgearth.org/en/latest/developer/shader_composition.html
     *
     * VirtualProgram (VP) is an osg::StateAttribute. But unlike most attributes, a VP
     * will inherit properties from other VPs in the state stack.
     *
     * Though the code has evolved quite a bit, VirtualProgram was originally adapted
     * from the VirtualProgram shader composition work done by Wojciech Lewandowski and
     * found in OSG's osgvirtualprogram example.
     */
    class OSGEARTH_EXPORT VirtualProgram : public osg::StateAttribute
    {
    public:
        static const osg::StateAttribute::Type SA_TYPE;

        /**
        * Gets the VP on a stateset, creating and installing one if the stateset
        * does not already have one. This is a convenient patternt to use, since
        * the normal use-case is to add functions to an existing VP rather than
        * to replace it entirely.
        */
        static VirtualProgram* getOrCreate(osg::StateSet* on);

        /**
        * Gets the VP on a stateset, or NULL if one is not found.
        */
        static VirtualProgram* get(osg::StateSet* on);
        static const VirtualProgram* get(const osg::StateSet* on);

        /**
        * Creates a new VP on the stateset, cloning an existing one if found.
        */
        static VirtualProgram* cloneOrCreate(osg::StateSet* stateset);

        /**
        * Creates a new VP on the "dest" stateset, either by cloning one found
        * on the "src" stateset, or otherwise just constructing a new one.
        */
        static VirtualProgram* cloneOrCreate(const osg::StateSet* src, osg::StateSet* dest);

    public:
        /**
         * Adds a custom shader function to the program.
         *
         * Call this method (rather than setShader directly) to inject "user" functions into the
         * shader program.
         *
         * name:      Name of the function. This should be the actual function name in the shader source.
         * source:    The shader source code.
         * location:  Function location relative to the built-ins.
         * order:     Lets you control the order of functions that you inject at the same location.
         *            The default order is 1.0. Note that many of osgEarth's built-in shaders (like
         *            those that render the terrain) use order=0.0 so that by default they run before
         *            user-injected functions by default.
         */
        void setFunction( 
            const std::string&           name,
            const std::string&           source, 
            ShaderComp::FunctionLocation location,
            float                        order =1.0f );

        void setFunction( 
            const std::string&           name,
            const std::string&           source, 
            ShaderComp::FunctionLocation location,
            ShaderComp::AcceptCallback*  acceptCallback,
            float                        order =1.0f );

        /** @deprecated and NOOP (remove after 2.10) */
        void setFunctionMinRange(const std::string& name, float minRange);

        /** @deprecated and NOOP (remove after 2.10) */
        void setFunctionMaxRange(const std::string& name, float maxRange);

        /**
         * Whether this VP should inherit shaders from parent state sets. This is
         * the normal operation. You can set this to "false" to "reset" the VP.
         */
        void setInheritShaders( bool value );
        bool getInheritShaders() const { return _inherit; }

        /**
        * Ability to add an extension to the VP
        */
        bool addGLSLExtension(const std::string& extension);
        bool hasGLSLExtension(const std::string& extension) const;
        bool removeGLSLExtension(const std::string& extension);
    public: 
        /**
         * Constructs a new VP
         */
        VirtualProgram( unsigned int mask = 0xFFFFFFFFUL );

        /**
         * Copy constructor
         */
        VirtualProgram( const VirtualProgram& VirtualProgram, 
                        const osg::CopyOp& copyop=osg::CopyOp::SHALLOW_COPY );

        META_StateAttribute( osgEarth, VirtualProgram, SA_TYPE);

        /** dtor */
        virtual ~VirtualProgram();

        /** 
         * Compare this program against another (used for state-sorting)
         * return -1 if *this < *rhs, 0 if *this==*rhs, 1 if *this>*rhs.
         */
        virtual int compare(const StateAttribute& sa) const;

        /**
         * If enabled, activate our program in the GL pipeline,
         * performing any rebuild operations that might be pending.
         */
        virtual void apply(osg::State& state) const;

        /**
         * Gets a shader by its ID.
         */
        PolyShader* getPolyShader( const std::string& shaderID ) const;

        /** 
         * Adds a shader to this VP's shader table.
         */
        osg::Shader* setShader( 
            const std::string&                 shaderID, 
            osg::Shader*                       shader,
            osg::StateAttribute::OverrideValue ov         =osg::StateAttribute::ON );
        
        osg::Shader* setShader(
            osg::Shader*                       shader,
            osg::StateAttribute::OverrideValue ov         =osg::StateAttribute::ON );

        /** Removes a shader from the local VP. */
        void removeShader( const std::string& shaderID );

        /** Add an attribute location binding. */
        void addBindAttribLocation( const std::string& name, GLuint index );

        /** Remove an attribute location binding. */
        void removeBindAttribLocation( const std::string& name );

        /** Gets a reference to the attribute bindings. */
        typedef osg::Program::AttribBindingList AttribBindingList;
        const AttribBindingList& getAttribBindingList() const { return _attribBindingList; }

        /** Access to the property template. */
        osg::Program* getTemplate() { return _template.get(); }
        const osg::Program* getTemplate() const { return _template.get(); }

        /** Enable logging to the console, for debugging. */
        void setShaderLogging( bool log );

        /** Enable logging to a file, for debugging. */
        void setShaderLogging( bool log, const std::string& filepath );

        /** Gets whether the accept callbacks vary per frame
          * @deprecated (remove after user support */
        bool getAcceptCallbacksVaryPerFrame() const;

        /** Sets whether the accept callbacks vary per frame
          * @deprecated (remove after user support) */
        void setAcceptCallbacksVaryPerFrame(bool acceptCallbacksVaryPerFrame);

        /** Inheritance mask */
        void setMask(unsigned mask) { _mask = mask; }
        unsigned getMask() const { return _mask; }

        /** Sets whether this VP is "abstract", i.e. pure virtual, and cannot be compiled on its own */
        void setIsAbstract(bool value) { _isAbstract = value; }
        bool getIsAbstract() const { return _isAbstract; }

    public: // StateAttribute
        virtual void compileGLObjects(osg::State& state) const;
        virtual void resizeGLObjectBuffers(unsigned maxSize);

        /** If State is non-zero, this function releases any associated OpenGL objects for
           * the specified graphics context. Otherwise, releases OpenGL objects
           * for all graphics contexts. */
        virtual void releaseGLObjects(osg::State* pState) const;

    public:
        typedef std::vector< osg::ref_ptr<osg::Shader> > ShaderVector;

    public:        
        struct ShaderEntry
        {
            ShaderEntry();
            bool accept(const osg::State& state) const;
            osg::ref_ptr<PolyShader>                 _shader;
            osg::StateAttribute::OverrideValue       _overrideValue;
            osg::ref_ptr<ShaderComp::AcceptCallback> _accept; // @deprecated
            bool operator < (const ShaderEntry& rhs) const;
        };

        typedef unsigned                              ShaderID;
        typedef vector_map<ShaderID, ShaderEntry>     ShaderMap;
        typedef std::set<std::string>                 ExtensionsSet;
        typedef std::map< std::string, std::string >  AttribAliasMap;
        typedef std::pair< std::string, std::string > AttribAlias;
        typedef std::vector< AttribAlias >            AttribAliasVector;

        //typedef osgEarth::fast_map< ProgramKey, ProgramEntry > ProgramMap;
        typedef std::pair< const osg::StateAttribute*, osg::StateAttribute::OverrideValue > AttributePair;
        typedef std::vector< AttributePair > AttrStack;

    public:
        /**
         * Populates the output collection with all the osg::Shader objects that
         * are applicable for the given State.
         * Returns the size of the output collection.
         */
        static int getShaders(
            const osg::State&                         state,
            std::vector<osg::ref_ptr<osg::Shader> >&  output);
        
        /**
         * Populates the output collection with all the PolyShaders that are applicable
         * for the given state. Information about the stage mask is not returned (TODO?)
         * Returns the size of the output collection.
         */
        static int getPolyShaders(
            const osg::State&                        state,
            std::vector<osg::ref_ptr<PolyShader> >&  output);

    public: // INTERNAL USE ONLY

        // thread-safe functions map getter
        void getFunctions( ShaderComp::FunctionLocationMap& out ) const;

        // thread-safe shader map copy
        void getShaderMap( ShaderMap& out ) const;

        // thread-safe shader accumulator
        void addShadersToAccumulationMap(VirtualProgram::ShaderMap& accumMap, const osg::State& state) const;

    protected: // serializable members

        // holds "template" data that should be installed in every auto-generated
        // Program, like uniform buffer bindings, etc.
        osg::ref_ptr<osg::Program> _template;

        unsigned int       _mask;
        AttribBindingList  _attribBindingList;

        // holds the injection points the user has added to this VP.
        // _dataModelMutex protects access to this member.
        ShaderComp::FunctionLocationMap _functions;

        // whether this VP inherits its state
        bool _inherit;

    protected: // non-serializable members

        // holds a map of each named shader installed in this VP.
        // _dataModelMutex protects access to this member.
        ShaderMap _shaderMap;

        ExtensionsSet _globalExtensions;

        // per-context cached shader map for thread-safe reuse without constant reallocation.
        struct ApplyVars
        {
            ShaderMap         accumShaderMap;
            ProgramKey        programKey;
            AttribBindingList accumAttribBindings;
            AttribAliasMap    accumAttribAliases;
        };
        mutable osg::buffered_object<ApplyVars> _apply;

        // protects access to the data members, which may be accessed by other VPs in the state stack.
        mutable Threading::Mutex _dataModelMutex;

        // The program cache holds an osg::Program instance for each collection of shaders
        // that comprises this VP. There can be multiple programs in the cache if the VP is
        // shared in the scene graph.
        mutable optional<bool> _active;
        bool _inheritSet;

        bool _logShaders;
        std::string _logPath;

        bool _isAbstract;

        UID _id;

        AttribAliasMap _attribAliases;

        bool _acceptCallbacksVaryPerFrame;

        mutable osg::buffered_object< osg::ref_ptr<osg::Program> > _lastUsedProgram;

        // Mechnism for remembering whether a VP has been applied during the same frame
        // and with the same attribute stack.
        struct AttrStackMemory
        {
            void remember(const osg::State& state, const AttrStack& rhs, osg::Program* p);
            osg::Program* recall(const osg::State& state, const AttrStack& rhs);

            struct Item
            {
                int frameNumber;
                AttrStack attrStack;
                osg::ref_ptr<osg::Program> program;
            };
            osg::buffered_object<Item> _item;
        };
        mutable AttrStackMemory _vpStackMemory;

        bool hasLocalFunctions() const;

        void accumulateFunctions(            
            const osg::State&                state,
            ShaderComp::FunctionLocationMap& result ) const;

        const AttribAliasMap& getAttribAliases() const { return _attribAliases; }

        // utility function
        static void accumulateShaders(
            const osg::State&  state,
            unsigned           mask,
            ShaderMap&         accumShaderMap,
            AttribBindingList& accumAttribBindings,
            AttribAliasMap&    accumAttribAliases,
            bool&              acceptCallbacksPresent);

        bool checkSharing();
    };

} // namespace osgEarth

#endif // OSGEARTH_VIRTUAL_PROGRAM_H
