//-----------------------------------------------------------------------------
// Copyright (c) 2012 GarageGames, LLC
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to
// deal in the Software without restriction, including without limitation the
// rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
// sell copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// 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.
//-----------------------------------------------------------------------------

#ifndef _REFBASE_H_
#define _REFBASE_H_

#ifndef _PLATFORMASSERT_H_
#  include "platform/platformAssert.h"
#endif
#ifndef _TYPETRAITS_H_
#  include "platform/typetraits.h"
#endif


/// Base class for objects which can be weakly referenced
/// (i.e., reference goes away when object is destroyed).
class WeakRefBase
{
public:

   /// Weak reference to WeakRefBase. 
   class WeakReference
   {
   public:

      [[nodiscard]] constexpr WeakRefBase* get() const { return mObject; }
      [[nodiscard]] constexpr U32 getRefCount()  const { return mRefCount; }

      constexpr void incRefCount() { mRefCount++; }
      constexpr void decRefCount() {
         AssertFatal( mRefCount > 0, "WeakReference - decrementing count of zero!" );
         if (--mRefCount==0)
            delete this;
      }
   private:

      friend class WeakRefBase;
      constexpr explicit WeakReference(WeakRefBase *object) :mObject(object), mRefCount(0) {}

      ~WeakReference() { AssertFatal(mObject==nullptr, "Deleting weak reference which still points at an object."); }

      // Object we reference
      WeakRefBase *mObject;

      // reference count for this structure (not WeakObjectRef itself)
      U32 mRefCount; 
   };

public:
   constexpr WeakRefBase() : mReference(nullptr) {}
   virtual ~WeakRefBase()  { clearWeakReferences(); }

   WeakReference* getWeakReference();

protected:
   void clearWeakReferences();

private:
   WeakReference * mReference;
};

template< typename T > class SimObjectPtr;

/// Weak reference pointer class.
/// Instances of this template class can be used as pointers to
/// instances of WeakRefBase and its subclasses.
/// When the object referenced by a WeakRefPtr instance is deleted,
/// the pointer to the object is set to NULL in the WeakRefPtr instance.
template <class T> class WeakRefPtr
{
public:
   constexpr WeakRefPtr()  : mReference(nullptr) {}
   WeakRefPtr(T *ptr)      : mReference(nullptr) { set(ptr); }
   WeakRefPtr(const WeakRefPtr<T> & ref) { mReference = nullptr; set(ref.mReference); }
   
   ~WeakRefPtr() { set(static_cast<WeakRefBase::WeakReference *>(nullptr)); }

   WeakRefPtr<T>& operator=(const WeakRefPtr<T>& ref)
   {
      if (this == &ref) { return *this; } // handle self assignment ( x = x; )
      set(ref.mReference);
      return *this;
   }
   WeakRefPtr<T>& operator=(T *ptr)
   {
      set(ptr);
      return *this;
   }

   /// Returns true if the pointer is not set.
   [[nodiscard]] constexpr bool isNull()        const { return mReference == nullptr || mReference->get() == nullptr; }
   
   /// Returns true if the pointer is set.
   [[nodiscard]] constexpr bool isValid()       const { return mReference && mReference->get(); }
   
   [[nodiscard]] constexpr T* operator->()      const { return getPointer(); }
   [[nodiscard]] constexpr T& operator*()       const { return *getPointer(); }
   [[nodiscard]] constexpr operator T*()        const { return getPointer(); }
   
   /// Returns the pointer.
   [[nodiscard]] constexpr T* getPointer()      const { return mReference ? (T*)mReference->get() : nullptr; }

protected:
   void set(WeakRefBase::WeakReference* ref)
   {
      if (mReference)
         mReference->decRefCount();
      mReference = nullptr;
      if (ref)
      {
         mReference = ref;
         mReference->incRefCount();
      }
   }

   void set(T* obj) { set(obj ? obj->getWeakReference() : nullptr); }
private:
   template< typename > friend class SimObjectPtr;
   WeakRefBase::WeakReference * mReference {nullptr};
};

/// Union of an arbitrary type with a WeakRefBase.  The exposed type will
/// remain accessible so long as the WeakRefBase is not cleared.  Once
/// it is cleared, accessing the exposed type will result in a NULL pointer.
template<class ExposedType>
class WeakRefUnion
{
   typedef WeakRefUnion<ExposedType> Union;

public:
   constexpr WeakRefUnion() : mPtr(nullptr) {}
   constexpr WeakRefUnion(const WeakRefPtr<WeakRefBase> & ref, ExposedType * ptr) : mPtr(nullptr) { set(ref, ptr); }
   constexpr WeakRefUnion(const Union & lock)                                     : mPtr(nullptr) { *this = lock; }
   constexpr WeakRefUnion(WeakRefBase * ref)       : mPtr(nullptr) { set(ref, dynamic_cast<ExposedType*>(ref)); }
   ~WeakRefUnion() { mWeakReference = nullptr; }

   Union & operator=(const Union & ptr)
   {
      if (this == *ptr) { return *this; }
      set(ptr.mWeakReference, ptr.getPointer());
      return *this;
   }

   void set(const WeakRefPtr<WeakRefBase> & ref, ExposedType * ptr)
   {
      mWeakReference = ref;
      mPtr = ptr;
   }

   [[nodiscard]] constexpr bool operator == (const ExposedType * t ) const { return getPointer() == t; }
   [[nodiscard]] constexpr bool operator != (const ExposedType * t ) const { return getPointer() != t; }
   [[nodiscard]] constexpr bool operator == (const Union &t )        const { return getPointer() == t.getPointer(); }
   [[nodiscard]] constexpr bool operator != (const Union &t )        const { return getPointer() != t.getPointer(); }
   [[nodiscard]] constexpr bool isNull()                             const { return mWeakReference.isNull() || !mPtr; }

   [[nodiscard]] constexpr ExposedType* getPointer()        const { return !mWeakReference.isNull() ? mPtr : nullptr; }
   [[nodiscard]] constexpr ExposedType* operator->()        const { return getPointer(); }
   [[nodiscard]] constexpr ExposedType& operator*()         const { return *getPointer(); }
   [[nodiscard]] constexpr operator ExposedType*()          const { return getPointer(); }

   [[nodiscard]] WeakRefPtr<WeakRefBase> getRef()           const { return mWeakReference; }

private:
   WeakRefPtr<WeakRefBase> mWeakReference;
   ExposedType * mPtr;
};

/// Base class for objects which can be strongly referenced
/// (i.e., as long as reference exists, object will exist,
/// when all strong references go away, object is destroyed).
class StrongRefBase : public WeakRefBase
{
   friend class StrongObjectRef;

public:
   StrongRefBase() { mRefCount = 0; }

   U32 getRefCount() const { return mRefCount; }

   /// object destroy self call (from StrongRefPtr).  Override if this class has specially allocated memory.
   virtual void destroySelf() { delete this; }

   /// Increments the reference count.
   void incRefCount()
   {
      mRefCount++;
   }

   /// Decrements the reference count.
   void decRefCount()
   {
      AssertFatal(mRefCount, "Decrementing a reference with refcount 0!");
      if(!--mRefCount)
         destroySelf();
   }

protected:
   U32 mRefCount; ///< reference counter for StrongRefPtr objects
};

/// Base class for StrongRefBase strong reference pointers.
class StrongObjectRef
{

public:
   /// Constructor, assigns from the object and increments its reference count if it's not NULL
   StrongObjectRef(StrongRefBase *object = nullptr) : mObject( object ) { incRef(); }

   /// Destructor, dereferences the object, if there is one
   ~StrongObjectRef() { decRef(); }

   /// Assigns this reference object from an existing StrongRefBase instance
   void set(StrongRefBase *object)
   {
      if( mObject != object )
      {
         decRef();
         mObject = object;
         incRef();
      }
   }

protected:
   StrongRefBase *mObject; ///< the object this RefObjectRef references

   /// increments the reference count on the referenced object
   constexpr void incRef() const
   {
      if(mObject)
         mObject->incRefCount();
   }

   /// decrements the reference count on the referenced object
   constexpr void decRef() const
   {
      if(mObject)
         mObject->decRefCount();
   }
};

/// Reference counted object template pointer class
/// Instances of this template class can be used as pointers to
/// instances of StrongRefBase and its subclasses.  The object will not
/// be deleted until all of the StrongRefPtr instances pointing to it
/// have been destructed.
template <class T> class StrongRefPtr : protected StrongObjectRef
{
public:
   StrongRefPtr() : StrongObjectRef() {}
   StrongRefPtr(T *ptr) : StrongObjectRef(ptr) {}
   StrongRefPtr(const StrongRefPtr<T>& ref) : StrongObjectRef(ref.mObject) {}
   ~StrongRefPtr() = default;

   StrongRefPtr<T>& operator=(const StrongRefPtr<T>& ref)
   {
      set(ref.mObject);
      return *this;
   }
   StrongRefPtr<T>& operator=(T *ptr)
   {
      set(ptr);
      return *this;
   }

   [[nodiscard]] constexpr bool isNull()  const  { return mObject == nullptr; }
   [[nodiscard]] constexpr bool isValid() const  { return mObject != nullptr; }
   [[nodiscard]] T* operator->() const { return getPointer(); }
   T& operator*() const  { return *getPointer(); }
   operator T*() const { return getPointer(); }
   T* getPointer() const { return const_cast<T*>(static_cast<T* const>(mObject)); }
};

/// Union of an arbitrary type with a StrongRefBase.  StrongRefBase will remain locked
/// until the union is destructed.  Handy for when the exposed class will
/// become invalid whenever the reference becomes invalid and you want to make sure that doesn't
/// happen.
template<class ExposedType>
class StrongRefUnion
{
   typedef StrongRefUnion<ExposedType> Union;

public:
   StrongRefUnion() : mPtr(nullptr) {}

   StrongRefUnion(const StrongRefPtr<StrongRefBase> & ref, ExposedType * ptr) : mPtr(nullptr) { set(ref, ptr); }
   StrongRefUnion(const Union & lock)                                         : mPtr(nullptr) { *this = lock; }
   StrongRefUnion(StrongRefBase * ref)         : mPtr(nullptr) { set(ref, dynamic_cast<ExposedType*>(ref)); }

   ~StrongRefUnion() { mReference = nullptr; }

   Union & operator=(const Union & ptr)
   {
      set(ptr.mReference, ptr.mPtr);
      return *this;
   }

   void set(const StrongRefPtr<StrongRefBase> & ref, ExposedType * ptr)
   {
      mReference = ref;
      mPtr = ptr;
   }

   [[nodiscard]] constexpr bool operator == (const ExposedType * t ) const { return mPtr == t; }
   [[nodiscard]] constexpr bool operator != (const ExposedType * t ) const { return mPtr != t; }
   [[nodiscard]] constexpr bool operator == (const Union &t )        const { return mPtr == t.mPtr; }
   [[nodiscard]] constexpr bool operator != (const Union &t )        const { return mPtr != t.mPtr; }
   [[nodiscard]] constexpr bool isNull()        const { return mReference.isNull() || !mPtr; }
   [[nodiscard]] constexpr bool isValid()       const { return mReference.isValid() && mPtr; }

   ExposedType* operator->() const { return mPtr; }
   ExposedType& operator*() const { return *mPtr; }
   operator ExposedType*() const { return mPtr; }
   ExposedType* getPointer() const { return mPtr; }

   StrongRefPtr<StrongRefBase> getRef() const { return mReference; }

private:
   StrongRefPtr<StrongRefBase> mReference;
   ExposedType * mPtr;
};

/// This oxymoron is a pointer that reference-counts the referenced
/// object but also NULLs out if the object goes away.
///
/// This is useful for situations where an object's lifetime is ultimately
/// governed by a superior entity but where individual objects may also die
/// independently of the superior entity.  All client code should use
/// StrongWeakRefs that keep object live as long as the superior entity doesn't
/// step in and kill them (in which case, the client code sees the reference
/// disappear).
template< class T >
class StrongWeakRefPtr
{
public:
   constexpr StrongWeakRefPtr()           : mReference( nullptr ) {}
   constexpr StrongWeakRefPtr( T* ptr )   : mReference( nullptr ) { _set( ptr ); }
   ~StrongWeakRefPtr()
   {
      if( mReference )
      {
         T* ptr = _get();
         if( ptr )
            ptr->decRefCount();

         mReference->decRefCount();
      }
   }

   [[nodiscard]] constexpr bool isNull()              const { return ( _get() == nullptr ); }
   [[nodiscard]] constexpr bool operator ==( T* ptr ) const { return ( _get() == ptr ); }
   [[nodiscard]] constexpr bool operator !=( T* ptr ) const { return ( _get() != ptr ); }
   [[nodiscard]] constexpr bool operator !()          const { return isNull(); }
   [[nodiscard]] constexpr T* operator ->()           const { return _get(); }
   [[nodiscard]] constexpr T& operator *()            const { return *( _get() ); }

   constexpr operator T*() const { return _get(); } // consider making this explicit

   T* getPointer() const { return _get(); }

   StrongWeakRefPtr& operator =( T* ptr )
   {
      _set( ptr );
      return *this;
   }

private:
   WeakRefBase::WeakReference* mReference;

   T* _get() const
   {
      if( mReference )
         return static_cast< T* >( mReference->get() );
      else
         return nullptr;
   }
   void _set( T* ptr )
   {
      if( mReference )
      {
         T* old = _get();
         if( old )
            old->decRefCount();

         mReference->decRefCount();
      }

      if( ptr )
      {
         ptr->incRefCount();
         mReference = ptr->getWeakReference();
         mReference->incRefCount();
      }
      else
         mReference = nullptr;
   }
};

//---------------------------------------------------------------

inline void WeakRefBase::clearWeakReferences()
{
   if (mReference)
   {
      mReference->mObject = nullptr;
      mReference->decRefCount();
      mReference = nullptr;
   }
}

inline WeakRefBase::WeakReference* WeakRefBase::getWeakReference()
{
   if (!mReference)
   {
      mReference = new WeakReference(this);
      mReference->incRefCount();
   }
   return mReference;
}

//---------------------------------------------------------------

template< typename T >
struct TypeTraits< WeakRefPtr< T > > : public _TypeTraits< WeakRefPtr< T > >
{
   typedef typename TypeTraits< T >::BaseType BaseType;
};
template< typename T >
struct TypeTraits< StrongRefPtr< T > > : public _TypeTraits< StrongRefPtr< T > >
{
   typedef typename TypeTraits< T >::BaseType BaseType;
};
template< typename T >
struct TypeTraits< StrongWeakRefPtr< T > > : public _TypeTraits< StrongWeakRefPtr< T > >
{
   typedef typename TypeTraits< T >::BaseType BaseType;
};

template< typename T >
[[nodiscard]] constexpr T& Deref( WeakRefPtr< T >& ref )
{
   return *ref;
}
template< typename T >
[[nodiscard]] constexpr T& Deref( StrongRefPtr< T >& ref )
{
   return *ref;
}
template< typename T >
[[nodiscard]] constexpr T& Deref( StrongWeakRefPtr< T >& ref )
{
   return *ref;
}

#endif