/* -*-c++-*- */
/* osgEarth - Geospatial SDK for OpenSceneGraph
* Copyright 2020 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/Threading>
#include <osgEarth/Containers>
#include <osgEarth/optional>
#include <osg/Shader>
#include <osg/Program>
#include <osg/StateAttribute>
#include <osg/buffered_value>
#include <string>
#include <map>

namespace osgEarth
{
    class VirtualProgram;

    namespace Util
    {
        class OSGEARTH_EXPORT ProgramRepo : public Threading::Mutexed<osg::Referenced>
        {
        public:
            // Program key (must be sorted set)
            using Key = std::set<unsigned>;

            // Referenced, b/c more than one ProgramKey can point to the same Entry
            // (shared program)
            struct Entry
            {
                using Ptr = std::shared_ptr<Entry>;
                osg::ref_ptr<osg::Program> _program;
                unsigned _frameLastUsed;
                std::unordered_set<UID> _users;
            };
            using ProgramMap = std::map<Key, Entry::Ptr>;

            //! 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 Key& key, unsigned frameNumber, UID user);

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

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

            //! Whether to automatically delete a Program after its final user releases it
            //! Defaults to true
            void setReleaseUnusedPrograms(bool value);

            //! Filesystem location in which to cache precopmiled shader binaries
            void setProgramBinaryCacheLocation(const std::string& directory);

            bool isProgramBinaryCachingActive() const;

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

            void resizeGLObjectBuffers(unsigned maxSize);

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

            //! Link a program, attempting to read/write its binary from the cache
            //! if one is specified
            void linkProgram(
                const Key&,
                osg::Program*, 
                osg::Program::PerContextProgram*,
                osg::State&);

            ProgramRepo();

            ~ProgramRepo();

            ProgramMap copy() const {
                lock();
                auto result = _db;
                unlock();
                return result;
            }

            //! Repository for simple, reusable VP objects.
            osg::ref_ptr<VirtualProgram> getOrCreateVirtualProgram(
                const std::string& name,
                std::function<VirtualProgram*()> create);

        private:
            mutable ProgramMap _db;
            bool _releaseUnusedPrograms;
            std::string _programBinaryCacheFolder;
            std::unordered_map<std::string, osg::ref_ptr<VirtualProgram>> _virtualProgramLUT;
        };
    }
}

namespace osgEarth
{
    using namespace osgEarth::Util;

    /**
     * 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.
     *
     * DO NOT DERIVE FROM VirtualProgram! IT WILL NOT WORK.
     */
    class OSGEARTH_EXPORT VirtualProgram final : public osg::StateAttribute
    {
    public:
        // User function injection points.
        enum FunctionLocation
        {
            // custom function to transform from model to view space
            LOCATION_VERTEX_TRANSFORM_MODEL_TO_VIEW,

            // 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);
            }
        };

        using OrderedFunction = std::pair<float, Function>;

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

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

        // 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
        };
        using StageMask = unsigned;

        /**
         * 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(VirtualProgram::FunctionLocation loc);
            VirtualProgram::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(unsigned mask) const;

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

            //! Is this component a fragment or output shader?
            bool isFragmentStage() const { return _nominalShader->getType() == osg::Shader::FRAGMENT; }

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

            //! Unique hash for this object
            unsigned getHash();

            /** 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,
                VirtualProgram::FunctionLocation location);

            static void clearShaderCache();

        protected:
            virtual ~PolyShader() { }

            std::string _name;
            std::string _source;
            VirtualProgram::FunctionLocation _location;
            optional<unsigned> _hash;

            // 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 std::mutex _cacheMutex;
            using PolyShaderCache = std::map<
                std::pair<std::string, std::string>,
                osg::ref_ptr<PolyShader>>;
            static PolyShaderCache _polyShaderCache;
        };

    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);

        /**
        * Whether to automatically delete a Program after its final user releases it
        * Defaults to true
        */
        static void setReleaseUnusedPrograms(bool value);

        /**
        * Folder to used to cache pre-compiled shader binaries
        */
        static void setProgramBinaryCacheLocation(const std::string& directory);

    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,
            FunctionLocation location,
            float order =1.0f );

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

        /**
         * 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);

        static void enableGLDebugging();

    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<AcceptCallback> _accept; // @deprecated

            // not used when using a vector_map, but keep around anyway
            bool operator < (const ShaderEntry& rhs) const;
        };

        using ShaderID = unsigned;
        // vector_map benchmarks much faster than std::map or std::unordered_map in this use case
        using ShaderMap = vector_map<ShaderID, ShaderEntry>;
        using ExtensionsSet = vector_set<std::string>;
        using AttribAliasMap = vector_map<std::string, std::string>;
        using AttribAlias = std::pair<std::string, std::string>;
        using AttribAliasVector = std::vector<AttribAlias>;
        using AttributePair = std::pair<const osg::StateAttribute*, osg::StateAttribute::OverrideValue>;
        using AttrStack = std::vector<AttributePair>;

    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.
         * 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(FunctionLocationMap& out) const;

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

        // thread-safe shader accumulator
        void addShadersToAccumulationMap(
            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.
        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;
            ProgramRepo::Key programKey;
            AttribBindingList accumAttribBindings;
            AttribAliasMap accumAttribAliases;
            ExtensionsSet accumExtensions;
        };
        mutable osg::buffered_object<ApplyVars> _apply;

        // protects access to the data members, which may be accessed by other VPs in the state stack.
        mutable std::mutex _dataModelMutex;
        bool _useDataModelMutex = true;

        // 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;

        static bool _gldebug;

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

        void accumulateFunctions(
            const osg::State& state,
            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,
            ExtensionsSet&     accumExtensions,
            bool&              acceptCallbacksPresent);

        bool checkSharing();

        //! Safely compare this VP to another VP.
        int compare_safe(const VirtualProgram& rhs) const;

    public: // backwards-compability functions

        inline void setFunction(
            const std::string& name,
            const std::string& source,
            unsigned location,
            float order = 1.0f)
        {
            setFunction(name, source, (FunctionLocation)location, order);
        }

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

    // For backwards compatibility
    namespace ShaderComp
    {
        using FunctionLocation = osgEarth::VirtualProgram::FunctionLocation;
        using AcceptCallback = osgEarth::VirtualProgram::AcceptCallback;
        using Function = osgEarth::VirtualProgram::Function;
        using OrderedFunction = osgEarth::VirtualProgram::OrderedFunction;
        using OrderedFunctionMap = osgEarth::VirtualProgram::OrderedFunctionMap;
        using FunctionLocationMap = osgEarth::VirtualProgram::FunctionLocationMap;
        using StageMask = osgEarth::VirtualProgram::StageMask;

        enum StageMaskValues {
            STAGE_VERTEX = osgEarth::VirtualProgram::STAGE_VERTEX,
            STAGE_TESSCONTROL = osgEarth::VirtualProgram::STAGE_TESSCONTROL,
            STAGE_TESSEVALUATION = osgEarth::VirtualProgram::STAGE_TESSEVALUATION,
            STAGE_GEOMETRY = osgEarth::VirtualProgram::STAGE_GEOMETRY,
            STAGE_FRAGMENT = osgEarth::VirtualProgram::STAGE_FRAGMENT,
            STAGE_COMPUTE = osgEarth::VirtualProgram::STAGE_COMPUTE
        };

        enum FunctionLocationValues {
            LOCATION_VERTEX_TRANSFORM_MODEL_TO_VIEW = osgEarth::VirtualProgram::LOCATION_VERTEX_TRANSFORM_MODEL_TO_VIEW,
            LOCATION_VERTEX_MODEL = osgEarth::VirtualProgram::LOCATION_VERTEX_MODEL,
            LOCATION_VERTEX_VIEW = osgEarth::VirtualProgram::LOCATION_VERTEX_VIEW,
            LOCATION_VERTEX_CLIP = osgEarth::VirtualProgram::LOCATION_VERTEX_CLIP,
            LOCATION_TESS_CONTROL = osgEarth::VirtualProgram::LOCATION_TESS_CONTROL,
            LOCATION_TESS_EVALUATION = osgEarth::VirtualProgram::LOCATION_TESS_EVALUATION,
            LOCATION_GEOMETRY = osgEarth::VirtualProgram::LOCATION_GEOMETRY,
            LOCATION_FRAGMENT_COLORING = osgEarth::VirtualProgram::LOCATION_FRAGMENT_COLORING,
            LOCATION_FRAGMENT_LIGHTING = osgEarth::VirtualProgram::LOCATION_FRAGMENT_LIGHTING,
            LOCATION_FRAGMENT_OUTPUT = osgEarth::VirtualProgram::LOCATION_FRAGMENT_OUTPUT,
            LOCATION_UNDEFINED = osgEarth::VirtualProgram::LOCATION_UNDEFINED
        };
    }

} // namespace osgEarth

#endif // OSGEARTH_VIRTUAL_PROGRAM_H