/* osgEarth
* Copyright 2025 Pelican Mapping
* MIT License
*/
#pragma once

#include <osgEarth/Common>
#include <osgEarth/GLUtils>
#include <osgEarth/TextureArena>
#include <osgEarth/Utils>
#include <osg/Geometry>
#include <osgUtil/RenderLeaf>
#include <osgUtil/RenderBin>
#include <vector>
#include <unordered_set>
#include <unordered_map>

// Usage hint to signal which textures are external materials. Currently only two slots are supported [0, 1]
#define CHONK_HINT_EXTENDED_MATERIAL_SLOT "osgEarth.Chonk.ExtendedMaterial.Slot"

namespace osgEarth
{
    class OSGEARTH_EXPORT ChonkMaterial
    {
    public:
        using Ptr = std::shared_ptr<ChonkMaterial>;

        static Ptr create() {
            return Ptr(new ChonkMaterial);
        }

        // texturearena indexes of material textures
        int albedo_index = -1;
        int normal_index = -1;
        int pbr_index = -1;
        int material1_index = -1;
        int material2_index = -1;

        // store refs to the original textures when the
        // factory objects's texture arena is using auto-release.
        // otherwise they are null.
        Texture::Ptr albedo_tex;
        Texture::Ptr normal_tex;
        Texture::Ptr pbr_tex;
        Texture::Ptr material1_tex;
        Texture::Ptr material2_tex;

    private:
        //! hidden to force use of create()
        ChonkMaterial() = default;
    };

    class ChonkFactory;
    class ChonkDrawable;

    /**
     * A bindless rendering unit comprised of one VBO and one EBO.
     * A single Chonk usually represents one drawable "Object" in
     * the scene. It can include up to 2 LODs.
     */
    class OSGEARTH_EXPORT Chonk
    {
    public:
        using element_t = GLuint;
        using Ptr = std::shared_ptr<Chonk>;
        using WeakPtr = std::weak_ptr<Chonk>;

        //! Adds a node.
        bool add(
            osg::Node* node,
            ChonkFactory& factory);

        //! Adds a node with screen-space-error limits.
        bool add(
            osg::Node* node,
            float far_pixel_scale,
            float near_pixel_scale,
            ChonkFactory& factory);

        //! Local bounding box of this Chonk.
        const osg::BoundingBoxf& getBound();

        //! Name, which will appear in GL debug labels
        std::string& name() { return _name; }
        const std::string& name() const { return _name; }

    public:
        //! Single vertex data.
        struct VertexGPU
        {
            osg::Vec3f position;
            osg::Vec3f normal;
            char normal_technique = (char)0;
            osg::Vec4ub color;
            osg::Vec2f uv;
            osg::Vec3f flex;

            // texturearena indexes of materials
            // TODO: use a material index instead?
            GLshort albedo_index = -1;
            GLshort normalmap_index = -1;
            GLshort pbr_index = -1;
            osg::Vec2s extended_material_index = osg::Vec2s(-1, -1);
        };

        static GLubyte NORMAL_TECHNIQUE_DEFAULT;
        static GLubyte NORMAL_TECHNIQUE_ZAXIS;
        static GLubyte NORMAL_TECHNIQUE_HEMISPHERE;

        static unsigned MATERIAL_VERTEX_SLOT;

        //! Geometry variant for a given pixel size range
        struct LOD
        {
            unsigned offset;
            std::size_t length;
            float far_pixel_scale;
            float near_pixel_scale;
        };

    public:

        //! Creates a new empty chonk.
        static Ptr create();

        std::string _name;
        std::vector<ChonkMaterial::Ptr> _materials;
        std::vector<VertexGPU> _vbo_store;
        std::vector<element_t> _ebo_store;
        std::vector<LOD> _lods;
        osg::BoundingBoxf _box;

        // Customed draw command
        using DrawCommand = DrawElementsIndirectBindlessCommandNV;
        using DrawCommands = std::vector<DrawCommand>;

        //! GL objects that can be used across shared GCs
        //! (because they are static and they are bindless)
        struct GLObjects : public BindlessShareableGLObjects
        {
            GLBuffer::Ptr vbo;
            GLBuffer::Ptr ebo;
            DrawCommands commands;
        };
        mutable osg::buffered_object<GLObjects> _globjects;

        //! Gets (or creates) a draw command, possibly
        //! resolving texture handles and uploading buffers.
        const DrawCommands& getOrCreateCommands(osg::State&) const;

    private:
        Chonk();

        friend class ChonkFactory;
        friend class ChonkDrawable;
    };

    /**
     * Converts OSG geometry into Chonk data.
     */
    class OSGEARTH_EXPORT ChonkFactory
    {
    public:
        using GetOrCreateFunction =
            std::function<Texture::Ptr(osg::Texture* osgTex, bool& isNew)>;

        //! Texture arena into which this factory will place textures it finds
        osg::ref_ptr<TextureArena> textures;

        //! Texture caching function
        GetOrCreateFunction getOrCreateTexture;

    public:
        //! Default, user should supply a TextureArena.
        ChonkFactory();

        //! Construct with a texture arena.
        ChonkFactory(TextureArena* textures);

        //! Function that can try to find textures instead of
        //! creating new ones. Good for sharing data across invocations, and
        //! for preventint duplicate textures in the arena.
        void setGetOrCreateFunction(GetOrCreateFunction);

        //! For the given node, this method populates the provided Chonk wit
        //! its geometry. If the factory finds any InstancedGeometry's along
        //! the way it will create a Chonk for each one in out_instances.
        //! Returns false if the node contained no geometry.
        bool load(
            osg::Node* node,
            Chonk* chonk, 
            float far_pixel_scale,
            float near_pixel_scale);

        bool load(
            osg::Node* node,
            ChonkDrawable* drawable,
            float far_pixel_scale,
            float near_pixel_scale);

        //! Stock "getOrCreateTexture" function for the ChonkFactory that works on
        //! a vector of Weak pointers. Use this or provide your own or don't use
        //! one at all.
        static GetOrCreateFunction getWeakTextureCacheFunction(
            std::vector<Texture::WeakPtr>& cache,
            std::mutex& cache_mutex);

    private:
        mutable std::vector<Texture::WeakPtr> _texcache;
        mutable std::mutex _texcache_mutex;
    };

    /**
     * Renders batches of chonks with gpu culling.
     */
    class OSGEARTH_EXPORT ChonkDrawable : public osg::Geometry //public osg::Drawable
    {
    public:
        META_Node(osgEarth, ChonkDrawable);

        using Vector = std::vector<osg::ref_ptr<ChonkDrawable>>;

    public:
        //! Construct a default drawable
        //! @param renderBinNumber OSG render bin number for the drawable
        ChonkDrawable(int render_bin_number = 3);

        //! Adds a node to the drawable. This method will rip all the geometry
        //! from the node, find any instances, and assemble everything into
        //! this drawable.
        bool add(
            osg::Node* node,
            ChonkFactory& factory,
            float far_pixel_scale = 1.0f,
            float near_pixel_scale = FLT_MAX);

        //! Adds one instance of a chonk to the drawable.
        void add(Chonk::Ptr value);

        //! Adds an instance of a chonk to the drawable.
        void add(
            Chonk::Ptr value,
            const osg::Matrixf& xform);

        //! Adds a chonk instance to the drawable, along with
        //! a tile-local UV coordinate. You only need this if you
        //! intend to set a custom culling shader.
        void add(
            Chonk::Ptr value,
            const osg::Matrixf& xform,
            const osg::Vec2f& local_uv);

        //! The timestamp at which this drawable was first created.
        //! Used for new data fade-in.
        void setBirthday(double value);
        double getBirthday() const;

        //! Set the camera distances for distance-based fading.
        void setFadeNearFar(float nearDist, float farDist);

        //! Sets the cutoff value for transparency
        void setAlphaCutoff(float value);

        //! Whether to GPU cull on the per-chonk basis.
        //! Default is true.
        void setUseGPUCulling(bool value);

        //! Render bin number to use for this drawable
        void setRenderBinNumber(int value);
        int getRenderBinNumber() const;

        //! Installs a basic chonk-rendering shader on a drawable.
        static void installRenderBin(ChonkDrawable*);

        //! Whether this drawable contains any geometry
        bool empty() const {
            return _batches.empty();
        }

    public:

        osg::BoundingBox computeBoundingBox() const override;
        osg::BoundingSphere computeBound() const override;
        void drawImplementation(osg::RenderInfo& ri) const override;
        void compileGLObjects(osg::RenderInfo& ri) const override;
        void resizeGLObjectBuffers(unsigned) override;
        void releaseGLObjects(osg::State*) const override;

    protected:
        virtual ~ChonkDrawable();

        mutable Mutex _m;

        // Keep me 16-byte aligned
        struct Instance
        {
            osg::Matrixf xform; // local xform
            osg::Vec2f uv; // tile-local UV
            GLuint lod;
            float visibility[2]; // per LOD
            float radius; // per chonk
            float alphaCutoff;
            GLint first_lod_cmd_index = -1; // invalid instance
        };
        using Instances = std::vector<Instance>;
        using Batches = std::unordered_map<Chonk::Ptr, Instances>;
        Batches _batches;
        bool _gpucull = true;
        float _fadeNear = 0.0f, _fadeFar = 0.0f;
        double _birthday = -1.0;
        float _alphaCutoff = 0.0f;
        int _renderBinNumber = 0;

        //! GL objects (and supporting data structures) that must
        //! be stored and managed for each unique OSG state,
        //! since they will differ for each view based on the camera
        //! location and will potentially change every frame.
        struct GLObjects : public PerStateGLObjects
        {
            // GPU variant/lod def. Keep me 16-byte aligned.
            struct ChonkLOD
            {
                osg::Vec3f center;
                GLfloat radius;
                GLfloat far_pixel_scale;
                GLfloat near_pixel_scale;
                GLfloat alpha_cutoff;
                GLfloat birthday;
                GLfloat fade_near;
                GLfloat fade_far;
                GLuint num_lods; // chonk-global
                GLuint total_num_commands; // global
            };

            // re-usable arrays
            std::vector<Chonk::DrawCommand> _commands;
            std::vector<Instance> _all_instances;
            std::vector<ChonkLOD> _chonk_lods;

            std::size_t _numInstances;
            std::size_t _maxNumLODs;
            osg::GLExtensions* _ext;
            GLVAO::Ptr _vao;
            GLBuffer::Ptr _commandBuf;
            GLBuffer::Ptr _instanceInputBuf;
            GLBuffer::Ptr _instanceOutputBuf;
            GLBuffer::Ptr _chonkBuf;
            bool _dirty = true;
            bool _gpucull = true;

            void(GL_APIENTRY * _glMultiDrawElementsIndirectBindlessNV)
                (GLenum, GLenum, const GLvoid*, GLsizei, GLsizei, GLint);

            struct PCPState {
                GLint _passUL = -1;
            };
            mutable std::unordered_map<const void*, PCPState> _pcps;

            void initialize(
                const osg::Object* host, osg::State& state);
            void update(
                const Batches&, const osg::Object* host,
                float fadeNear, float fadeFar,
                double birthday, float alphaCutoff,
                osg::State&);
            void cull(osg::State& state);
            void draw(osg::State& state);
            void release();
        };
        mutable osg::buffered_object<GLObjects> _globjects;

        void update_and_cull_batches(osg::State& state) const;

        void draw_batches(osg::State& state) const;

        void dirtyGLObjects();

    public:
        // support intersectors and stats visitors
        bool supports(const osg::PrimitiveFunctor& f) const override { return true; }
        void accept(osg::PrimitiveFunctor&) const override;

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

    private:
        mutable bool _proxy_dirty = true;
        mutable osg::Vec3Array* _proxy_verts = nullptr;
        mutable osg::DrawElementsUInt* _proxy_indices = nullptr;
        void refreshProxy() const;
        void update_gl_objects(GLObjects&, osg::State&) const;
        ChonkDrawable(const ChonkDrawable& rhs, const osg::CopyOp&) { }

        friend class ChonkRenderBin;
    };

    /**
     * Custom renderbin for rendering ChonkDrawables en masse
     * Called "ChonkBin"
     */
    class OSGEARTH_EXPORT ChonkRenderBin : public osgUtil::RenderBin
    {
    public:
        ChonkRenderBin();
        ChonkRenderBin(const ChonkRenderBin& rhs, const osg::CopyOp& op);

        void drawImplementation(
            osg::RenderInfo& renderInfo,
            osgUtil::RenderLeaf*&) override;

        virtual osg::Object* cloneType() const override { return new ChonkRenderBin(); }
        virtual osg::Object* clone(const osg::CopyOp& copyop) const override { return new ChonkRenderBin(*this, copyop); } // note only implements a clone of type.
        virtual bool isSameKindAs(const osg::Object* obj) const override { return dynamic_cast<const ChonkRenderBin*>(obj) != 0L; }
        virtual const char* libraryName() const override { return "osgEarth"; }
        virtual const char* className() const override { return "ChonkRenderBin"; }

        static void releaseSharedGLObjects(osg::State* state);

    private:
        osg::ref_ptr<osg::StateSet> _cullSS;

        friend class ChonkDrawable;

        osg::ref_ptr<osgUtil::StateGraph> _cull_sg;

        struct CullLeaf : public Util::CustomRenderLeaf {
            CullLeaf(osgUtil::RenderLeaf*);
            void draw(osg::State&) override;
        };
        struct DrawLeaf : public Util::CustomRenderLeaf {
            DrawLeaf(osgUtil::RenderLeaf*, bool, bool);
            void draw(osg::State&) override;
            bool _first, _last;
        };
    };
}
