/* osgEarth
 * Copyright 2025 Pelican Mapping
 * MIT License
 */

#ifndef OSGEARTH_DRAW_INSTANCED_H
#define OSGEARTH_DRAW_INSTANCED_H 1

#include <osgEarth/Common>
#include <osgEarth/Containers>
#include <osgEarth/VirtualProgram>
#include <osg/NodeVisitor>
#include <osg/Drawable>
#include <osg/Geometry>
#include <osg/TextureBuffer>

/**
 * Some utilities to support *DrawInstanced rendering.
 */
namespace osgEarth { namespace Util
{
    using namespace osgEarth;

    namespace DrawInstanced
    {
        class OSGEARTH_EXPORT InstanceGeometry : public osg::Geometry
        {
        public:
            InstanceGeometry();
            InstanceGeometry(const osg::Geometry& geometry);
            InstanceGeometry(const InstanceGeometry&, const osg::CopyOp& copyop = osg::CopyOp::SHALLOW_COPY);
            META_Node(osgEarth::Util::DrawInstanced, InstanceGeometry);

            //! Positioning matrices for each instance. 
            //! NOTE that each matrix includes an embedded object ID - do not use it directly;
            //! call decodeMatrix() to extract it.
            const std::vector< osg::Matrixf >& getMatrices() const;
            void setMatrices(const std::vector< osg::Matrixf >& matrices);

            //! Decodes an encoded matrix into one you can use for math.
            static osg::Matrixf decodeMatrix(const osg::Matrixf& matrix);

            //! A mesh comprising all geometry instances; you can use this for intersection
            //! or bounds testing
            osg::Geometry* getProxyGeometry() const {
                return _mesh;
            }

        public: // osg::Geometry

            bool supports(const osg::PrimitiveFunctor& f) const override { return true; }
            void accept(osg::PrimitiveFunctor& f) const override;

            bool supports(const osg::PrimitiveIndexFunctor& f) const override { return true; }
            void accept(osg::PrimitiveIndexFunctor& f) const override;

            osg::BoundingBox computeBoundingBox() const override;

        private:
            std::vector<osg::Matrixf> _matrices;
            osg::ref_ptr<osg::Geometry> _mesh;
        };

        /**
         * Visitor that converts all the primitive sets in a graph to use
         * instanced draw calls.
         * Called by convertGraphToUseDrawInstanced().
         */
        class OSGEARTH_EXPORT ConvertToDrawInstanced : public osg::NodeVisitor
        {
        public:
            /**
             * Create the visitor that will convert primitive sets to draw
             * <num> instances.
             */
            ConvertToDrawInstanced(
                unsigned                numInstances,
                bool                    optimize,
                osg::TextureBuffer*     tbo,
                int                     firstTboUnit);

            void apply(osg::Drawable&);
            void apply(osg::LOD&);
            void apply(osg::Node&);

        public:
            int getTextureImageUnit() const { return _tboUnit; }

        protected:
            unsigned _numInstances;
            bool _optimize;
            std::list<osg::PrimitiveSet*> _primitiveSets;
            osg::TextureBuffer* _tbo;
            int _tboUnit;
            osg::ref_ptr<osg::Drawable::ComputeBoundingBoxCallback> _bboxComputer;
        };


        /**
         * Creates a virtual shader program that implements DrawInstanced rendering.
         * You should prepare the scene graph first by calling
         * convertGraphToUseDrawInstanced().
         * @return false If instancing is not available
         */
        extern OSGEARTH_EXPORT bool install(osg::StateSet* stateset);
        extern OSGEARTH_EXPORT void remove (osg::StateSet* stateset);


        /**
         * Processes a scene graph and converts all the top-level MatrixTransform
         * nodes into shader uniforms that can be used with the VirtualProgram
         * created by createDrawInstacedShaders.
         * NOTE: You must also call install(StateSet) to activate instancing.
         * @return false If instancing is not available
         */
        extern OSGEARTH_EXPORT bool convertGraphToUseDrawInstanced(
            osg::Group* graph );
    }
} }

#endif // OSGEARTH_DRAW_INSTANCED_H
