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

#include <osgEarth/Map>
#include <osgEarth/Terrain>
#include <osgEarth/TerrainOptions>
#include <osgEarth/TerrainEffect>
#include <osgEarth/TerrainTileModelFactory>
#include <osgEarth/TerrainTileNode>
#include <osgEarth/TerrainEngineRequirements>
#include <osgEarth/TerrainResources>
#include <osgEarth/ShaderUtils>
#include <osgEarth/Progress>
#include <osgEarth/TileKey>
#include <osg/CoordinateSystemNode>
#include <osg/Geode>
#include <osg/NodeCallback>
#include <osg/BoundingBox>
#include <osgUtil/RenderBin>
#include <set>

#define OSGEARTH_ENV_TERRAIN_ENGINE_DRIVER "OSGEARTH_TERRAIN_ENGINE"


namespace osgUtil {
    class CullVisitor;
}

namespace osgEarth
{
    class TerrainEffect;
    class VirtualProgram;

    namespace Util
    {
        class TerrainEngineNodeFactory;
    }

    /**
     * Terrain Engine interface. (Pure virtual)
     * Access to the low-level API for the terrain engine, the node graph that
     * renders the planet's surface.
     */
    class /*interface*/ TerrainEngine
    {
    public:
        /**
         * Callback for customizing the TileModel.
         */
        class CreateTileModelCallback : public osg::Referenced
        {
        public:
            virtual void onCreateTileModel(
                TerrainEngineNode* engine,
                TerrainTileModel*  model) = 0;
        };

        using ComputeTileRangeCallback = std::function<float(osg::Node* node, osg::NodeVisitor& nv)>;
        using ComputeTilePixelSizeCallback = std::function<float(osg::Node* node, osg::NodeVisitor& nv)>;

    public:
        //! Unique ID of the terrain engine instance
        virtual UID getUID() const = 0;

        //! Create a tile data model from the map
        virtual TerrainTileModel* createTileModel(
            const Map* map,
            const TileKey& key,
            const CreateTileManifest& manifest,
            ProgressCallback* progress) = 0;

        //! Access the shared GL resources associated with the terrain engine
        virtual TerrainResources* getResources() const = 0;

        //! Map associated with this terrain engine
        virtual const Map* getMap() const = 0;

        //! Gets the Terrain interface for interacting with the scene graph
        virtual Terrain* getTerrain() = 0;
        virtual const Terrain* getTerrain() const = 0;

        //! Adds a terrain effect
        virtual void addEffect(TerrainEffect* effect) = 0;

        //! Removes a terrain effect
        virtual void removeEffect(TerrainEffect* effect) = 0;

        //! Marks the terrain tiles intersecting the provied extent as invalid
        virtual void invalidateRegion(const GeoExtent& extent) = 0;

        //! Marks the terrain tiles intersecting the provied extent as invalid
        virtual void invalidateRegion(
            const GeoExtent& extent, 
            unsigned minLevel) = 0;

        //! Marks the terrain tiles intersecting the provied extent as invalid
        virtual void invalidateRegion(
            const GeoExtent& extent, 
            unsigned minLevel, 
            unsigned maxLevel) = 0;

        //! Marks the terrain tiles intersecting the provied extent as invalid
        virtual void invalidateRegion(
            const std::vector<const Layer*> layers,
            const GeoExtent& extent,
            unsigned minLevel,
            unsigned maxLevel) = 0;

        //! Marks the terrain tiles intersecting the provied extent as invalid
        virtual void invalidateRegion(
            const std::vector<const Layer*> layers,
            const GeoExtent& extent) = 0;

        //! Access the stateset used to render the entire terrain.
        virtual osg::StateSet* getTerrainStateSet() = 0;

        //! Access the stateset used to render terrain surface layers.
        virtual osg::StateSet* getSurfaceStateSet() = 0;

        //! Scene graph root of the terrain engine
        virtual osg::Node* getNode() = 0;

        //! Adds a TileModel creation callback, so you can add custom data
        //! to the TileModel after it's created.
        virtual void addCreateTileModelCallback(CreateTileModelCallback* callback) = 0;
        virtual void removeCreateTileModelCallback(CreateTileModelCallback* callback) = 0;

        //! Create a standalone terrain tile. This method is not used by the
        //! actual terrain engine, but rather exists to support external
        //! processes that need tile geometry. The terrain engine implementation
        //! may or may not provide this capability and will return NULL if
        //! it does not.
        //! @param model Tile model for which to build a tile
        //! @param flags Bitwise-OR of the CreateTileFlags enums (MAY BE IGNORED)
        //! @param referenceLOD If non-zero, adjust the vertex dimensions of the returned tile
        //!        to match this LOD. Example: ask for a tile at tileKey.lod=15, tile is 17x17.
        //!        Specific a referenceLOD=16, tile will be 33x33.
        //! @param area Restrict the processed area to a subarea of the tile model.
        //!        If not valid then the entire tile model will be used.
        virtual osg::Node* createStandaloneTile(
            const TerrainTileModel* model,
            int createTileFlags,
            unsigned referenceLOD,
            const TileKey& subRegion) = 0;

        //! Returns the name of the JobArena that handles terrain tile loading,
        //! or an empty string if the engine does not use an arena.
        virtual std::string getJobArenaName() const { return ""; }

        //! Access an immutable options API
        virtual TerrainOptionsAPI getOptions() = 0;

        //! Number of terrain tiles resident in memory 
        //! (including cached dormant tiles not being rendered)
        virtual unsigned getNumResidentTiles() const = 0;

        //! Tell the engine you updates options.
        virtual void dirtyTerrainOptions() = 0;

        //! Set a function that computes the on-screen size of a tile 
        //! for the purposes of LOD selection.
        virtual void setComputeTilePixelSizeCallback(const ComputeTilePixelSizeCallback& callback) = 0;

        //! Set a function that computes the range of a tile for the purposes of culling.
        virtual void setComputeTileRangeCallback(const ComputeTileRangeCallback& callback) = 0;
    };

    /**
     * TerrainEngineNode is the base class and interface for map engine implementations.
     *
     * A map engine lives under a MapNode and is responsible for generating the
     * actual geometry representing the Earth.
     */
    class OSGEARTH_EXPORT TerrainEngineNode : public osg::CoordinateSystemNode,  public TerrainEngine
    {
    public:
        TerrainTileModel* createTileModel(
            const Map* map,
            const TileKey& key,
            const CreateTileManifest& manifest,
            ProgressCallback* progress) override;

        const Map* getMap() const override { return _map.get(); }

        TerrainResources* getResources() const override;

        virtual const TerrainEngineRequirements& getRequirements() const = 0;

        Terrain* getTerrain() override { return _terrainInterface.get(); }
        const Terrain* getTerrain() const override { return _terrainInterface.get(); }

        void addEffect( TerrainEffect* effect ) override;

        void removeEffect( TerrainEffect* effect ) override;

        void invalidateRegion(const GeoExtent& extent) override {
            invalidateRegion(extent, 0, INT_MAX);
        }

        void invalidateRegion(const GeoExtent& extent, unsigned minLevel) override {
            invalidateRegion(extent, minLevel, INT_MAX);
        }

        virtual void invalidateRegion(const GeoExtent& extent, unsigned minLevel, unsigned maxLevel) override {
            //NOP by default
        }

        virtual void invalidateRegion(
            const std::vector<const Layer*> layers,
            const GeoExtent& extent,
            unsigned minLevel,
            unsigned maxLevel) override
        {
            //NOP by default
        }

        void invalidateRegion(
            const std::vector<const Layer*> layers,
            const GeoExtent& extent) override
        {
            invalidateRegion(layers, extent, 0, INT_MAX);
        }

        virtual osg::StateSet* getSurfaceStateSet() override { return getOrCreateStateSet(); }

        //! Adds a TileModel creation callback, so you can add custom data
        //! to the TileModel after it's created.
        void addCreateTileModelCallback(CreateTileModelCallback* callback) override;
        void removeCreateTileModelCallback(CreateTileModelCallback* callback) override;

        //! Flags supporting the createStandaloneTile method
        enum CreateTileFlags
        {
            CREATE_TILE_INCLUDE_TILES_WITH_CONSTRAINTS = 1,
            CREATE_TILE_INCLUDE_TILES_WITH_MASKS = 1, // backward compatibility
            CREATE_TILE_INCLUDE_TILES_WITHOUT_CONSTRAINTS = 2,
            CREATE_TILE_INCLUDE_TILES_WITHOUT_MASKS = 2, // backward compatibility
            CREATE_TILE_INCLUDE_ALL = ~0
        };

        //! Create a standalone terrain tile. This method is not used by the
        //! actual terrain engine, but rather exists to support external
        //! processes that need tile geometry. The terrain engine implementation
        //! may or may not provide this capability and will return NULL if
        //! it does not.
        //! @param model Tile model for which to build a tile
        //! @param flags Bitwise-OR of the CreateTileFlags enums (MAY BE IGNORED)
        //! @param referenceLOD If non-zero, adjust the vertex dimensions of the returned tile
        //!        to match this LOD. Example: ask for a tile at tileKey.lod=15, tile is 17x17.
        //!        Specific a referenceLOD=16, tile will be 33x33.
        //! @param area Restrict the processed area to a subarea of the tile model.
        //!        If not valid then the entire tile model will be used.
        virtual osg::Node* createStandaloneTile(
            const TerrainTileModel* model,
            int createTileFlags,
            unsigned referenceLOD,
            const TileKey& subRegion) { return nullptr; }

        //! Shut down the engine
        virtual void shutdown();

        void setComputeTileRangeCallback(const ComputeTileRangeCallback& callback) override { _computeTileRangeCallback = callback; }

        void setComputeTilePixelSizeCallback(const ComputeTilePixelSizeCallback& callback) override { _computeTilePixelSizeCallback = callback; }
        
        ComputeTileRangeCallback& getComputeTileRangeCallback() { return _computeTileRangeCallback; }
        ComputeTilePixelSizeCallback& getComputeTilePixelSizeCallback() { return _computeTilePixelSizeCallback; }

        //! Return the top level node associated with the terrain.
        osg::Node* getNode() override {
            return this;
        }

        //! API for accessing or changing terrain options.
        //! You should call dirtyTerrainOptions() after changing the options.
        TerrainOptionsAPI getOptions() override {
            return TerrainOptionsAPI(&_optionsConcrete);
        }

        //! Subclass may override this to do something when the terrain options change
        virtual void dirtyTerrainOptions() { }

    public:
        class OSGEARTH_EXPORT ModifyTileBoundingBoxCallback : public osg::Referenced
        {
        public:
            virtual void modifyBoundingBox(const TileKey& key, osg::BoundingBox& box) const =0;
        };

        void addModifyTileBoundingBoxCallback(ModifyTileBoundingBoxCallback* callback);
        void removeModifyTileBoundingBoxCallback(ModifyTileBoundingBoxCallback* callback);

    public:

        static TerrainEngineNode* create(const TerrainOptions& options);

    protected:
        TerrainEngineNode();
        virtual ~TerrainEngineNode();

    public: // osg::Node overrides
        virtual osg::BoundingSphere computeBound() const;
        virtual void traverse( osg::NodeVisitor& );

    protected:
        friend class MapNode;
        friend class TerrainEngineNodeFactory;

        //! Assigns a map to render
        void setMap(const Map* map, const TerrainOptions& options);

        //! Subclass runs this after a call to setMap.
        virtual void onSetMap() { }

        // signals that a redraw is needed because something changed.
        virtual void requestRedraw();

        // Request the state setup be refreshed because something has changed that requires new
        // shaders, state, etc.
        virtual void dirtyState() { }

        osg::ref_ptr<TerrainResources> _textureResourceTracker;
        osg::ref_ptr<const Map> _map;

    private:
        friend struct TerrainEngineNodeCallbackProxy;

        void onMapModelChanged( const MapModelChange& change );
        virtual void updateTextureCombining() { }

    protected:
        osg::ref_ptr<Terrain> _terrainInterface;
        unsigned _dirtyCount;
        bool _updateScheduled;
        TerrainOptions _optionsConcrete;

        typedef std::vector<osg::ref_ptr<TerrainEffect> > TerrainEffectVector;
        TerrainEffectVector effects_;

        typedef std::vector<osg::ref_ptr<CreateTileModelCallback> > CreateTileModelCallbacks;
        CreateTileModelCallbacks _createTileModelCallbacks;
        mutable Threading::ReadWriteMutex _createTileModelCallbacksMutex;

        typedef std::vector<osg::ref_ptr<ModifyTileBoundingBoxCallback> > ModifyTileBoundingBoxCallbacks;
        ModifyTileBoundingBoxCallbacks _modifyTileBoundingBoxCallbacks;

        osg::ref_ptr<TerrainTileModelFactory> _tileModelFactory;

        ComputeTileRangeCallback _computeTileRangeCallback;
        ComputeTilePixelSizeCallback _computeTilePixelSizeCallback;

    public:
        // internal
        void fireModifyTileBoundingBoxCallbacks(const TileKey& key, osg::BoundingBox& box);

        /** Access a typed effect. */
        template<typename T>
        T* getEffect() {
            for(TerrainEffectVector::iterator i = effects_.begin(); i != effects_.end(); ++i ) {
                T* e = dynamic_cast<T*>(i->get());
                if ( e ) return e;
            }
            return 0L;
        }
    };
} // namespace osgEarth
