/* osgEarth
* Copyright 2025 Pelican Mapping
* MIT License
*/
#ifndef OSGEARTH_UTILS_H
#define OSGEARTH_UTILS_H 1

#include <osgEarth/Common>
#include <osgEarth/StringUtils>
#include <osgEarth/Threading>
#include <osgEarth/DrawInstanced>

#include <osg/Vec3f>
#include <osg/AutoTransform>
#include <osg/OperationThread>
#include <osg/TriangleIndexFunctor>
#include <osg/MatrixTransform>
#include <osgGA/GUIEventHandler>
#include <osgViewer/View>
#include <osgUtil/CullVisitor>
#include <osgUtil/RenderBin>
#include <osgUtil/RenderLeaf>

#include <string>
#include <list>
#include <map>
#include <typeinfo>
#include <stack>

namespace osgEarth { namespace Util
{
    using namespace osgEarth::Threading;

    /**
    * Tracking object - place this in a user data container to track the
    * construction and destruction of objects.
    */
    struct TrackerTag : public osg::Object
    {
        META_Object(osgEarth, TrackerTag);

        std::shared_ptr<std::atomic_int> _count;

        TrackerTag() : _count(nullptr) { }
        TrackerTag(const TrackerTag& rhs, const osg::CopyOp&) : _count(rhs._count) { }
        TrackerTag(std::shared_ptr<std::atomic_int> count) : _count(count) {
            ++(*_count);
        }
        ~TrackerTag() {
            if (_count) --(*_count);
        }
    };

    /**
    * General purpose wrapper that wraps a value in an osg::Object so you can
    * easily store it in a user data container (or whatever)
    */
    template<class T>
    struct WrapperObject : public osg::Object
    {
        META_Object(osgEarth, WrapperObject<T>);
        WrapperObject() { }
        WrapperObject(const std::string name, const T& t) : value(t) { setName(name); }
        WrapperObject(const WrapperObject<T>& rhs, const osg::CopyOp& op) : osg::Object(rhs, op), value(rhs.value) { }
        T value;
    };

    /**
     * Utility for storing observables in an osg::Object
     */
    class ObjectStorage
    {
    private:
        template<typename T>
        class Data : public osg::Object {
        public:
            META_Object(osgEarth, Data);
            Data(const std::string& name, T* obj) : _obj(obj) { setName(name); }
            Data(const std::string& name, std::shared_ptr<T> obj) : _shared_obj(obj) { setName(name); }
            Data() { }
            Data(const Data& rhs, const osg::CopyOp& copy) : osg::Object(rhs, copy), _obj(rhs._obj) { }
            osg::observer_ptr<T> _obj;
            std::shared_ptr<T> _shared_obj;
        };

    public:
        template<typename T>
        static void set(osg::Object* o, const std::string& name, T* obj) {
            if (o == nullptr || obj == nullptr) return;
            osg::UserDataContainer* udc = o->getOrCreateUserDataContainer();
            unsigned i = udc->getUserObjectIndex(name);
            if (i < udc->getNumUserObjects()) udc->removeUserObject(i);
            udc->addUserObject(new Data<T>(name, obj));
        }
        template<typename T>
        static bool get(const osg::Object* o, const std::string& name, osg::ref_ptr<T>& out) {
            out = nullptr;
            if (o == nullptr) return false;
            const Data<T>* data = dynamic_cast<const Data<T>*>(osg::getUserObject(o, name));
            return data ? data->_obj.lock(out) : false;
        }
        template<typename T>
        static void set(osg::Object* o, T* obj) {
            if (o == nullptr || obj == nullptr) return;
            const char* name = typeid(T).name();
            osg::UserDataContainer* udc = o->getOrCreateUserDataContainer();
            unsigned i = udc->getUserObjectIndex(name);
            if (i < udc->getNumUserObjects()) udc->removeUserObject(i);
            udc->addUserObject(new Data<T>(name, obj));
        }
        template<typename T>
        static bool get(const osg::Object* o, osg::ref_ptr<T>& out) {
            out = nullptr;
            if (o == nullptr) return false;
            const char* name = typeid(T).name();
            const Data<T>* data = dynamic_cast<const Data<T>*>(osg::getUserObject(o, name));
            return data ? data->_obj.lock(out) : false;
        }
        template<typename T>
        static void remove(osg::Object* o, T* obj) {
            const char* name = typeid(T).name();
            osg::UserDataContainer* udc = o->getUserDataContainer();
            if (udc) {
                unsigned i = udc->getUserObjectIndex(name);
                if (i < udc->getNumUserObjects()) udc->removeUserObject(i);
            }
        }
        template<typename T>
        static void set(osg::Object* o, std::shared_ptr<T> obj) {
            if (o == nullptr || obj == nullptr) return;
            const char* name = typeid(T).name();
            osg::UserDataContainer* udc = o->getOrCreateUserDataContainer();
            unsigned i = udc->getUserObjectIndex(name);
            if (i < udc->getNumUserObjects()) udc->removeUserObject(i);
            udc->addUserObject(new Data<T>(name, obj));
        }
        template<typename T>
        static bool get(const osg::Object* o, std::shared_ptr<T>& out) {
            out = nullptr;
            const char* name = typeid(T).name();
            if (!o) return false;
            const Data<T>* data = dynamic_cast<const Data<T>*>(osg::getUserObject(o, name));
            out = data->_shared_obj;
            return data->_shared_obj != nullptr;
        }

        template<typename T>
        struct Install : public osg::NodeCallback
        {
            osg::ref_ptr<T> _data;

            Install(T* data) : _data(data) { }

            void operator()(osg::Node* node, osg::NodeVisitor* nv)
            {
                ObjectStorage::set(nv, _data.get());
                traverse(node, nv);
            }
        };

        template<typename T>
        struct SetValue : public osg::NodeCallback
        {
            std::string _key;
            T _value;

            SetValue(const std::string& key, const T& value) : _key(key), _value(value) { }
            void operator()(osg::Node* node, osg::NodeVisitor* nv)
            {
                nv->setUserValue(_key, _value);
                traverse(node, nv);
            }
        };
    };

    struct Utils
    {
        /**
         * Clamps v to [vmin..vmax], then remaps its range to [r0..r1].
         */
        static double remap( double v, double vmin, double vmax, double r0, double r1 )
        {
            float vr = (osg::clampBetween(v, vmin, vmax)-vmin)/(vmax-vmin);
            return r0 + vr * (r1-r0);
        }
    };

    /**
     * Proxy class that registers a custom render bin's prototype with the
     * rendering system
     */
    template<class T>
    struct osgEarthRegisterRenderBinProxy
    {
        osgEarthRegisterRenderBinProxy(const std::string& name)
        {
            _prototype = new T();
            osgUtil::RenderBin::addRenderBinPrototype(name, _prototype.get());
        }

        ~osgEarthRegisterRenderBinProxy()
        {
            osgUtil::RenderBin::removeRenderBinPrototype(_prototype.get());
            _prototype = 0L;
        }

        osg::ref_ptr<T> _prototype;
    };

    struct OSGEARTH_EXPORT RenderBinUtils
    {
        static unsigned getTotalNumRenderLeaves(osgUtil::RenderBin* bin);
    };

    /**
     * RenderLeaf that copies another renderleaf's data and then provied
     * a custom draw function.
     */
    class OSGEARTH_EXPORT CustomRenderLeaf : public osgUtil::RenderLeaf
    {
    public:
        CustomRenderLeaf(osgUtil::RenderLeaf* leaf);

        //! Implemenet this for a custom draw
        virtual void draw(osg::State& state) = 0;

        // from RenderLeaf
        void render(osg::RenderInfo& renderInfo, osgUtil::RenderLeaf* previous) override;
    };

    /**
     * Shim to apply vertex cache optimizations to geometry when it's legal.
     * This is really only here to work around an OSG bug in the VertexAccessOrder
     * optimizer, which corrupts non-Triangle geometries.
     */
    struct OSGEARTH_EXPORT VertexCacheOptimizer : public osg::NodeVisitor
    {
        VertexCacheOptimizer();
        virtual ~VertexCacheOptimizer() { }
        void apply(osg::Drawable& drawable);
    };

    /**
     * Sets the data variance on all discovered drawables.
     */
    struct OSGEARTH_EXPORT SetDataVarianceVisitor : public osg::NodeVisitor
    {
        SetDataVarianceVisitor(osg::Object::DataVariance value);
        virtual ~SetDataVarianceVisitor() { }
        void apply(osg::Drawable& drawable);
        osg::Object::DataVariance _value;
    };

    /**
     * Scans geometry and validates that it's set up properly.
     */
    struct OSGEARTH_EXPORT GeometryValidator : public osg::NodeVisitor
    {
        GeometryValidator();
        void apply(osg::Group& group);
        void apply(osg::Geometry& geom);
    };

    /**
     * Allocates and merges buffer objects for each Drawable in the scene graph.
     */
    class OSGEARTH_EXPORT AllocateAndMergeBufferObjectsVisitor : public osg::NodeVisitor
    {
    public:
        AllocateAndMergeBufferObjectsVisitor();
        virtual ~AllocateAndMergeBufferObjectsVisitor() { }
        void apply(osg::Drawable& drawable);
    };

    /**
     * Adapts an osg::Callback (or derived class) to accept a lambda.
     */
    template<typename T = osg::Callback>
    struct LambdaCallback : public T {
        LambdaCallback(std::function<void(osg::NodeVisitor&)> func) : _func(func) { }
        std::function<void(osg::NodeVisitor&)> _func;
        bool run(osg::Object* object, osg::Object* data) override {
            _func(*reinterpret_cast<osg::NodeVisitor*>(data));
            return T::traverse(object, data);
        }
    };

    /**
     * Tracks usage data by maintaining a sentry-blocked linked list.
     * Each time a use called "use" the corresponding record moves to
     * the right of the sentry marker. After a cycle you can call
     * collectTrash to process all users that did not call use() in the
     * that cycle, and dispose of them.
     */
    template<typename T>
    class SentryTracker
    {
    public:
        struct ListEntry
        {
            ListEntry(T data, void* token) : _data(data), _token(token) { }
            T _data;
            void* _token;
        };

        using List = std::list<ListEntry>;
        using ListIterator = typename List::iterator;
        using Token = ListIterator;

        SentryTracker()
        {
            reset();
        }

        ~SentryTracker()
        {
            for (auto& e : _list)
            {
                Token* te = static_cast<Token*>(e._token);
                if (te)
                    delete te;
            }
        }

        void reset()
        {
            for (auto& e : _list)
            {
                Token* te = static_cast<Token*>(e._token);
                if (te)
                    delete te;
            }
            _list.clear();
            _list.emplace_front(nullptr, nullptr); // the sentry marker
            _sentryptr = _list.begin();
            _total = 0;
        }

        std::size_t size() const
        {
            return _total;
        }

        List _list;
        ListIterator _sentryptr;
        std::atomic_uint _total = { 0u };

        inline bool erase(void* token)
        {
            if (token == nullptr)
                return false;

            Token* ptr = static_cast<Token*>(token);
            _list.erase(*ptr);
            delete ptr;
            --_total;
            return true;
        }

        inline void* use(const T& data, void* token)
        {
            if (token)
            {
                Token* ptr = static_cast<Token*>(token);

                // Keep data current in case caller reuses a token for a replaced object.
                (*ptr)->_data = data;

                _list.splice(_list.begin(), _list, *ptr);
                *ptr = _list.begin();
                return ptr;
            }
            else
            {
                Token* ptr = new Token();
                _list.emplace_front(data, ptr);
                *ptr = _list.begin();
                ++_total;
                return ptr;
            }
        }

        template<class CALLABLE>
        inline void flush(unsigned maxCount, CALLABLE&& dispose)
        {
            // After cull, all visited tiles are in front of the sentry, and all
            // non-visited tiles are behind it. Start at the sentry position and
            // iterate over the non-visited tiles, checking them for deletion.
            ListIterator i = _sentryptr;
            ListIterator tmp;
            unsigned count = 0;

            for (++i; i != _list.end() && count < maxCount; ++i)
            {
                ListEntry& le = *i;

                // if we are ready to dispose it, dispose() will return true, and we can safely delete the entry.
                if (dispose(le._data))
                {
                    // back up the iterator so we can safely erase the entry:
                    tmp = i;
                    --i;

                    // delete the token
                    delete static_cast<Token*>(le._token);

                    // remove it from the tracker list:
                    _list.erase(tmp);
                    ++count;
                    --_total;
                }                
            }

            // reset the sentry.
            _list.splice(_list.begin(), _list, _sentryptr);
            _sentryptr = _list.begin();
        }
    };


    /** @deprecated - please use ObjectStorage instead */
    template<typename T>
    class OptionsData : public osg::Object {
    public:
        static void set(osg::Object* o, const std::string& name, T* obj) {
            osg::UserDataContainer* udc = o->getOrCreateUserDataContainer();
            unsigned i = udc->getUserObjectIndex(name);
            if (i == udc->getNumUserObjects()) udc->removeUserObject(i);
            udc->addUserObject(new OptionsData<T>(name, obj));
        }
        static bool lock(const osg::Object* o, const std::string& name, osg::ref_ptr<T>& out) {
            if (!o) return false;
            const OptionsData<T>* data = dynamic_cast<const OptionsData<T>*>(osg::getUserObject(o, name));
            return data ? data->_obj.lock(out) : false;
        }
        static osg::ref_ptr<T> get(const osg::Object* o, const std::string& name) {
            osg::ref_ptr<T> result;
            bool ok = lock(o, name, result);
            return result;
        }
    public:
        META_Object(osgEarth, OptionsData);
        OptionsData(const std::string& name, T* obj) : _obj(obj) { setName(name); }
        OptionsData() { }
        OptionsData(const OptionsData& rhs, const osg::CopyOp& copy) : osg::Object(rhs, copy), _obj(rhs._obj) { }
    private:
        osg::observer_ptr<T> _obj;
    };

    /**
     * An osg::Operation that runs a user function ONCE and then
     * disappears
     */
    struct OneTimer : public osg::Operation
    {
        using Function = std::function<void()>;

        OneTimer(Function func) :
            osg::Operation("osgEarth::OneTimer", true),
            _func(func) { }

        void operator()(osg::Object* obj) override
        {
            if (getKeep())
                _func();
            setKeep(false);
        }

    private:
        Function _func;
    };


    /**
     * Visitor that visits each instace of hte template node type.
     */
    template<typename T>
    struct TypedNodeVisitor : public osg::NodeVisitor
    {
        using Function = std::function<void(T&, const osg::Matrix&)>;
        Function _func;
        std::stack<osg::Matrix> _transformStack;

        TypedNodeVisitor(const Function& func) :
            osg::NodeVisitor(),
            _func(func)
        {
            setTraversalMode(TRAVERSE_ALL_CHILDREN);
            setNodeMaskOverride(~0);
            _transformStack.push(osg::Matrix::identity());
        }

        void apply(osg::Node& node) override
        {
            T* n = dynamic_cast<T*>(&node);
            if (n) _func(*n, _transformStack.top());
            traverse(node);
        }

        void apply(osg::Transform& node) override
        {
            osg::Matrix m = _transformStack.empty() ? osg::Matrix() : _transformStack.top();
            node.computeLocalToWorldMatrix(m, this);
            _transformStack.push(m);
            apply(static_cast<osg::Group&>(node));
            _transformStack.pop();
        }
    };

    /**
     * Visitor that visits each triangle in each geometry
     */
    struct TriangleVisitor : public osg::NodeVisitor
    {
    public:
        using Function = std::function<void(
            osg::Geometry& geom,
            unsigned i0,
            unsigned i1,
            unsigned i2,
            const osg::Matrix& l2w)>;

    public:
        TriangleVisitor(const Function& func) : 
            osg::NodeVisitor(),
            _func(func)
        {
            setTraversalMode(TRAVERSE_ALL_CHILDREN);
            setNodeMaskOverride(~0);
            _transformStack.push(osg::Matrix::identity());
        }

    private:
        struct TriangleFullSend {
            TriangleVisitor* _visitor;
            osg::Geometry* _geom;
            inline void operator () (unsigned i0, unsigned i1, unsigned i2)
            {
                _visitor->_func(*_geom, i0, i1, i2, _visitor->_transformStack.top());
            }
        };

    public: // overrides
        void apply(osg::Transform& node) override
        {
            osg::Matrix m = _transformStack.empty() ? osg::Matrix() : _transformStack.top();
            node.computeLocalToWorldMatrix(m, this);
            _transformStack.push(m);
            traverse(node);
            _transformStack.pop();
        }

        void apply(osg::Geometry& geom) override
        {
            auto instanced_geom = dynamic_cast<Util::DrawInstanced::InstanceGeometry*>(&geom);
            auto geom_to_use = instanced_geom ? instanced_geom->getProxyGeometry() : &geom;

            osg::TriangleIndexFunctor<TriangleFullSend> _sender;
            _sender._visitor = this;
            _sender._geom = geom_to_use;
            geom.accept(_sender);
        }

    private:
        Function _func;
        std::stack<osg::Matrix> _transformStack;
    };
} }

#endif // OSGEARTH_UTILS_H
