// SPDX-FileComment: This file is part of TNL - Template Numerical Library (https://tnl-project.org/)
// SPDX-License-Identifier: MIT

#pragma once

#include "SharedPointer.h"

#include <TNL/Assert.h>
#include <TNL/Backend.h>
#include <TNL/Allocators/Default.h>
#include <TNL/Devices/Cuda.h>
#include <TNL/Pointers/SmartPointer.h>
#include <TNL/Pointers/SmartPointersRegister.h>

#include <cstring>    // std::memcpy, std::memcmp
#include <cstddef>    // std::nullptr_t
#include <algorithm>  // swap

namespace TNL::Pointers {

/**
 * \brief Specialization of the \ref SharedPointer for the CUDA device.
 *
 * \tparam  Object is a type of object to be owned by the pointer.
 */
template< typename Object >
class SharedPointer< Object, Devices::Cuda > : public SmartPointer
{
private:
   /**
    * \typedef Enabler
    *
    * Convenient template alias for controlling the selection of copy- and
    * move-constructors and assignment operators using SFINAE.
    * The type Object_ is "enabled" iff Object_ and Object are not the same,
    * but after removing const and volatile qualifiers they are the same.
    */
   template< typename Object_ >
   using Enabler =
      std::enable_if_t< ! std::is_same_v< Object_, Object > && std::is_same_v< std::remove_cv_t< Object >, Object_ > >;

   // friend class will be needed for templated assignment operators
   template< typename Object_, typename Device_ >
   friend class SharedPointer;

public:
   /**
    * \brief Type of the object owned by the pointer.
    */
   using ObjectType = Object;

   /**
    * \brief Type of the device where the object is to be mirrored.
    */
   using DeviceType = Devices::Cuda;

   /**
    * \brief Type of the allocator for \e DeviceType.
    */
   using AllocatorType = typename Allocators::Default< DeviceType >::Allocator< ObjectType >;

   /**
    * \brief Constructor of empty pointer.
    */
   SharedPointer( std::nullptr_t )
   : pd( nullptr ),
     cuda_pointer( nullptr )
   {}

   /**
    * \brief Constructor with parameters of the Object constructor.
    *
    * \tparam Args is variadic template type of arguments of the Object constructor.
    * \tparam args are arguments passed to the Object constructor.
    */
   template< typename... Args >
   explicit SharedPointer( Args... args )
   : pd( nullptr ),
     cuda_pointer( nullptr )
   {
      this->allocate( args... );
   }

   /**
    * \brief Constructor with initializer list.
    *
    * \tparam Value is type of the initializer list elements.
    * \param list is the instance of the initializer list..
    */
   template< typename Value >
   explicit SharedPointer( std::initializer_list< Value > list )
   : pd( nullptr ),
     cuda_pointer( nullptr )
   {
      this->allocate( list );
   }

   /**
    * \brief Constructor with nested initializer lists.
    *
    * \tparam Value is type of the nested initializer list elements.
    * \param list is the instance of the nested initializer list..
    */
   template< typename Value >
   explicit SharedPointer( std::initializer_list< std::initializer_list< Value > > list )
   : pd( nullptr ),
     cuda_pointer( nullptr )
   {
      this->allocate( list );
   }

   /**
    * \brief Copy constructor.
    *
    * \param pointer is the source shared pointer.
    */
   SharedPointer( const SharedPointer& pointer )  // this is needed only to avoid the default compiler-generated constructor
   : pd( pointer.pd ),
     cuda_pointer( pointer.cuda_pointer )
   {
      this->pd->counter += 1;
   }

   /**
    * \brief Copy constructor.
    *
    * This is specialization for compatible object types.
    *
    * \param pointer is the source shared pointer.
    */
   template< typename Object_, typename = Enabler< Object_ > >
   SharedPointer( const SharedPointer< Object_, DeviceType >& pointer )  // conditional constructor for non-const -> const data
   : pd( const_cast< PointerData* >( pointer.pd ) ),
     cuda_pointer( pointer.cuda_pointer )
   {
      this->pd->counter += 1;
   }

   /**
    * \brief Move constructor.
    *
    * \param pointer is the source shared pointer.
    */
   SharedPointer( SharedPointer&& pointer ) noexcept  // this is needed only to avoid the default compiler-generated constructor
   : pd( pointer.pd ),
     cuda_pointer( pointer.cuda_pointer )
   {
      pointer.pd = nullptr;
      pointer.cuda_pointer = nullptr;
   }

   /**
    * \brief Move constructor.
    *
    * This is specialization for compatible object types.
    *
    * \param pointer is the source shared pointer.
    */
   template< typename Object_, typename = Enabler< Object_ > >
   SharedPointer( SharedPointer< Object_, DeviceType >&& pointer )  // conditional constructor for non-const -> const data
   : pd( const_cast< PointerData* >( pointer.pd ) ),
     cuda_pointer( pointer.cuda_pointer )
   {
      pointer.pd = nullptr;
      pointer.cuda_pointer = nullptr;
   }

   /**
    * \brief Create new object based in given constructor parameters.
    *
    * \tparam Args is variadic template type of arguments to be passed to the
    *    object constructor.
    * \param args are arguments to be passed to the object constructor.
    * \return true if recreation was successful, false otherwise.
    */
   template< typename... Args >
   bool
   recreate( Args... args )
   {
#ifdef TNL_DEBUG_SHARED_POINTERS
      std::cerr << "Recreating shared pointer to " << getType< ObjectType >() << '\n';
#endif
      if( ! this->pd )
         return this->allocate( args... );

      if( this->pd->counter == 1 ) {
         /****
          * The object is not shared -> recreate it in-place, without reallocation
          */
         this->pd->data.~Object();
         new( &this->pd->data ) Object( args... );
         Backend::memcpy(
            reinterpret_cast< void* >( this->cuda_pointer ),
            reinterpret_cast< const void* >( &this->pd->data ),
            sizeof( Object ),
            Backend::MemcpyHostToDevice );
         this->set_last_sync_state();
         return true;
      }

      // free will just decrement the counter
      this->free();

      return this->allocate( args... );
   }

   /**
    * \brief Arrow operator for accessing the object owned by constant smart pointer.
    *
    * \return constant pointer to the object owned by this smart pointer. It
    * returns pointer to object image on the CUDA device if it is called from CUDA
    * kernel and pointer to host image otherwise.
    */
   __cuda_callable__
   const Object*
   operator->() const
   {
      TNL_ASSERT_TRUE( this->pd, "Attempt to dereference a null pointer" );
#if defined( __CUDA_ARCH__ ) || defined( __HIP_DEVICE_COMPILE__ )
      return this->cuda_pointer;
#else
      return &this->pd->data;
#endif
   }

   /**
    * \brief Arrow operator for accessing the object owned by non-constant smart pointer.
    *
    * \return pointer to the object owned by this smart pointer. It
    * returns pointer to object image on the CUDA device if it is called from CUDA
    * kernel and pointer to host image otherwise.
    */
   __cuda_callable__
   Object*
   operator->()
   {
      TNL_ASSERT_TRUE( this->pd, "Attempt to dereference a null pointer" );
#if defined( __CUDA_ARCH__ ) || defined( __HIP_DEVICE_COMPILE__ )
      return this->cuda_pointer;
#else
      this->pd->maybe_modified = true;
      return &this->pd->data;
#endif
   }

   /**
    * \brief Dereferencing operator for accessing the object owned by constant smart pointer.
    *
    * \return constant reference to the object owned by this smart pointer. It
    * returns reference to object image on the CUDA device if it is called from CUDA
    * kernel and reference to host image otherwise.
    */
   __cuda_callable__
   const Object&
   operator*() const
   {
      TNL_ASSERT_TRUE( this->pd, "Attempt to dereference a null pointer" );
#if defined( __CUDA_ARCH__ ) || defined( __HIP_DEVICE_COMPILE__ )
      return *( this->cuda_pointer );
#else
      return this->pd->data;
#endif
   }

   /**
    * \brief Dereferencing operator for accessing the object owned by non-constant smart pointer.
    *
    * \return reference to the object owned by this smart pointer.  It
    * returns reference to object image on the CUDA device if it is called from CUDA
    * kernel and reference to host image otherwise.
    */
   __cuda_callable__
   Object&
   operator*()
   {
      TNL_ASSERT_TRUE( this->pd, "Attempt to dereference a null pointer" );
#if defined( __CUDA_ARCH__ ) || defined( __HIP_DEVICE_COMPILE__ )
      return *( this->cuda_pointer );
#else
      this->pd->maybe_modified = true;
      return this->pd->data;
#endif
   }

   /**
    * \brief Conversion to boolean type.
    *
    * \return Returns true if the pointer is not empty, false otherwise.
    */
   __cuda_callable__
   operator bool() const
   {
      return this->pd != nullptr;
   }

   /**
    * \brief Negation operator.
    *
    * \return Returns false if the pointer is not empty, true otherwise.
    */
   __cuda_callable__
   bool
   operator!() const
   {
      return this->pd == nullptr;
   }

   /**
    * \brief Constant object reference getter.
    *
    * No synchronization of this pointer will be performed due to calling
    * this method.
    *
    * \tparam Device says what image of the object one want to dereference. It
    * can be either \ref DeviceType or Devices::Host.
    * \return constant reference to the object image on given device.
    */
   template< typename Device = Devices::Host >
   [[nodiscard]] __cuda_callable__
   const Object&
   getData() const
   {
      static_assert(
         std::is_same_v< Device, Devices::Host > || std::is_same_v< Device, Devices::Cuda >,
         "Only Devices::Host or Devices::Cuda devices are accepted here." );
      TNL_ASSERT_TRUE( this->pd, "Attempt to dereference a null pointer" );
      TNL_ASSERT_TRUE( this->cuda_pointer, "Attempt to dereference a null pointer" );
      if( std::is_same_v< Device, Devices::Host > )
         return this->pd->data;
      if( std::is_same_v< Device, Devices::Cuda > )
         return *( this->cuda_pointer );
   }

   /**
    * \brief Non-constant object reference getter.
    *
    * After calling this method, the object owned by the pointer might need
    * to be synchronized. One should not forget to call
    * \ref Pointers::synchronizeSmartPointersOnDevice< Devices::Cuda >()
    * before calling CUDA kernel using object from this smart pointer.
    *
    * \tparam Device says what image of the object one want to dereference. It
    * can be either \ref DeviceType or Devices::Host.
    * \return constant reference to the object image on given device.
    */
   template< typename Device = Devices::Host >
   [[nodiscard]] __cuda_callable__
   Object&
   modifyData()
   {
      static_assert(
         std::is_same_v< Device, Devices::Host > || std::is_same_v< Device, Devices::Cuda >,
         "Only Devices::Host or Devices::Cuda devices are accepted here." );
      TNL_ASSERT_TRUE( this->pd, "Attempt to dereference a null pointer" );
      TNL_ASSERT_TRUE( this->cuda_pointer, "Attempt to dereference a null pointer" );
      if( std::is_same_v< Device, Devices::Host > ) {
         this->pd->maybe_modified = true;
         return this->pd->data;
      }
      if( std::is_same_v< Device, Devices::Cuda > )
         return *( this->cuda_pointer );
   }

   /**
    * \brief Assignment operator.
    *
    * It assigns object owned by the pointer \e ptr to \e this pointer.
    *
    * \param ptr input pointer
    * \return constant reference to \e this
    */
   const SharedPointer&
   operator=( const SharedPointer& ptr )  // this is needed only to avoid the default compiler-generated operator
   {
      this->free();
      this->pd = ptr.pd;
      this->cuda_pointer = ptr.cuda_pointer;
      if( this->pd != nullptr )
         this->pd->counter += 1;
#ifdef TNL_DEBUG_SHARED_POINTERS
      std::cerr << "Copy-assigned shared pointer: counter = " << this->pd->counter << ", type: " << getType< ObjectType >()
                << '\n';
#endif
      return *this;
   }

   /**
    * \brief Assignment operator for compatible object types.
    *
    * It assigns object owned by the pointer \e ptr to \e this pointer.
    *
    * \param ptr input pointer
    * \return constant reference to \e this
    */
   template< typename Object_, typename = Enabler< Object_ > >
   const SharedPointer&
   operator=( const SharedPointer< Object_, DeviceType >& ptr )  // conditional operator for non-const -> const data
   {
      this->free();
      this->pd = const_cast< PointerData* >( ptr.pd );
      this->cuda_pointer = ptr.cuda_pointer;
      if( this->pd != nullptr )
         this->pd->counter += 1;
#ifdef TNL_DEBUG_SHARED_POINTERS
      std::cerr << "Copy-assigned shared pointer: counter = " << this->pd->counter << ", type: " << getType< ObjectType >()
                << '\n';
#endif
      return *this;
   }

   /**
    * \brief Move operator.
    *
    * It assigns object owned by the pointer \e ptr to \e this pointer.
    *
    * \param ptr input pointer
    * \return constant reference to \e this
    */
   const SharedPointer&
   operator=( SharedPointer&& ptr ) noexcept  // this is needed only to avoid the default compiler-generated operator
   {
      this->free();
      this->pd = ptr.pd;
      this->cuda_pointer = ptr.cuda_pointer;
      ptr.pd = nullptr;
      ptr.cuda_pointer = nullptr;
#ifdef TNL_DEBUG_SHARED_POINTERS
      std::cerr << "Move-assigned shared pointer: counter = " << this->pd->counter << ", type: " << getType< ObjectType >()
                << '\n';
#endif
      return *this;
   }

   /**
    * \brief Move operator.
    *
    * It assigns object owned by the pointer \e ptr to \e this pointer.
    *
    * \param ptr input pointer
    * \return constant reference to \e this
    */
   template< typename Object_, typename = Enabler< Object_ > >
   const SharedPointer&
   operator=( SharedPointer< Object_, DeviceType >&& ptr )  // conditional operator for non-const -> const data
   {
      this->free();
      this->pd = const_cast< PointerData* >( ptr.pd );
      this->cuda_pointer = ptr.cuda_pointer;
      ptr.pd = nullptr;
      ptr.cuda_pointer = nullptr;
#ifdef TNL_DEBUG_SHARED_POINTERS
      std::cerr << "Move-assigned shared pointer: counter = " << this->pd->counter << ", type: " << getType< ObjectType >()
                << '\n';
#endif
      return *this;
   }

   /**
    * \brief Cross-device pointer synchronization.
    *
    * This method is usually called by the smart pointers register when calling
    * \ref Pointers::synchronizeSmartPointersOnDevice< Devices::Cuda >()
    *
    * \return true if the synchronization was successful, false otherwise.
    */
   bool
   synchronize() override
   {
      if( this->pd == nullptr )
         return true;
      if( this->modified() ) {
#ifdef TNL_DEBUG_SHARED_POINTERS
         std::cerr << "Synchronizing shared pointer: counter = " << this->pd->counter << ", type: " << getType< ObjectType >()
                   << '\n';
         std::cerr << "   ( " << sizeof( Object ) << " bytes, CUDA address " << this->cuda_pointer << " )\n";
#endif
         TNL_ASSERT_NE( this->cuda_pointer, nullptr, "" );
         Backend::memcpy( (void*) this->cuda_pointer, (void*) &this->pd->data, sizeof( Object ), Backend::MemcpyHostToDevice );
         this->set_last_sync_state();
         return true;
      }
      return true;
   }

   /**
    * \brief Reset the pointer to empty state.
    */
   void
   clear()
   {
      this->free();
   }

   /**
    * \brief Swap the owned object with another pointer.
    *
    * \param ptr2 the other shared pointer for swapping.
    */
   void
   swap( SharedPointer& ptr2 ) noexcept
   {
      std::swap( this->pd, ptr2.pd );
      std::swap( this->cuda_pointer, ptr2.cuda_pointer );
   }

   /**
    * \brief Destructor.
    */
   ~SharedPointer() override
   {
      this->free();
      getSmartPointersRegister< DeviceType >().remove( this );
   }

protected:
   struct PointerData
   {
      Object data;
      char data_image[ sizeof( Object ) ];
      int counter = 1;
      bool maybe_modified = true;

      template< typename... Args >
      explicit PointerData( Args... args )
      : data( args... )
      {}
   };

   template< typename... Args >
   bool
   allocate( Args... args )
   {
      this->pd = new PointerData( args... );
      // allocate on device
      this->cuda_pointer = AllocatorType{}.allocate( 1 );
      // synchronize
      this->synchronize();
#ifdef TNL_DEBUG_SHARED_POINTERS
      std::cerr << "Created shared pointer to " << getType< ObjectType >() << " (cuda_pointer = " << this->cuda_pointer
                << ")\n";
#endif
      getSmartPointersRegister< DeviceType >().insert( this );
      return true;
   }

   void
   set_last_sync_state()
   {
      TNL_ASSERT_TRUE( this->pd, "Attempt to dereference a null pointer" );
      std::memcpy( (void*) &this->pd->data_image, (void*) &this->pd->data, sizeof( Object ) );
      this->pd->maybe_modified = false;
   }

   bool
   modified()
   {
      TNL_ASSERT_TRUE( this->pd, "Attempt to dereference a null pointer" );
      // optimization: skip bitwise comparison if we're sure that the data is the same
      if( ! this->pd->maybe_modified )
         return false;
      return std::memcmp(
                reinterpret_cast< const void* >( &this->pd->data_image ),
                reinterpret_cast< const void* >( &this->pd->data ),
                sizeof( Object ) )
          != 0;
   }

   void
   free()
   {
      if( this->pd ) {
#ifdef TNL_DEBUG_SHARED_POINTERS
         std::cerr << "Freeing shared pointer: counter = " << this->pd->counter << ", cuda_pointer = " << this->cuda_pointer
                   << ", type: " << getType< ObjectType >() << '\n';
#endif
         if( ! --this->pd->counter ) {
            delete this->pd;
            this->pd = nullptr;
            if( this->cuda_pointer )
               AllocatorType{}.deallocate( this->cuda_pointer, 1 );
#ifdef TNL_DEBUG_SHARED_POINTERS
            std::cerr << "...deleted data.\n";
#endif
         }
      }
   }

   PointerData* pd;

   // cuda_pointer can't be part of PointerData structure, since we would be
   // unable to dereference this-pd on the device
   Object* cuda_pointer;
};

}  // namespace TNL::Pointers
