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

#pragma once

#include <unordered_set>
#include <unordered_map>

#include <TNL/Pointers/SmartPointer.h>
#include <TNL/Timer.h>
#include <TNL/Backend.h>
#include <TNL/Devices/Sequential.h>
#include <TNL/Devices/Host.h>

namespace TNL::Pointers {

/**
 * Since TNL currently supports only execution on host (which does not need
 * to register and synchronize smart pointers) and CUDA GPU's, the smart
 * pointers register is implemented only for CUDA. If more execution types
 * which need to register smart pointers are implemented in the future, this
 * should beome a class template specialization.
 */
class SmartPointersRegister
{
public:
   /**
    * Negative deviceId means that \ref TNL::Backend::getDevice
    * will be called to get the device ID.
    */
   void
   insert( SmartPointer* pointer, int deviceId = -1 )
   {
      if( deviceId < 0 )
         deviceId = Backend::getDevice();
      pointersOnDevices[ deviceId ].insert( pointer );
   }

   /**
    * Negative deviceId means that \ref TNL::Backend::getDevice
    * will be called to get the device ID.
    */
   void
   remove( SmartPointer* pointer, int deviceId = -1 )
   {
      if( deviceId < 0 )
         deviceId = Backend::getDevice();
      try {
         pointersOnDevices.at( deviceId ).erase( pointer );
      }
      catch( const std::out_of_range& ) {
         std::cerr << "Given deviceId " << deviceId << " does not have any pointers yet. "
                   << "Requested to remove pointer " << pointer << ". "
                   << "This is most likely a bug in the smart pointer.\n";
         throw;
      }
   }

   /**
    * Negative deviceId means that \ref TNL::Backend::getDevice
    * will be called to get the device ID.
    */
   bool
   synchronizeDevice( int deviceId = -1 )
   {
      if( deviceId < 0 )
         deviceId = Backend::getDevice();
      try {
         const auto& set = pointersOnDevices.at( deviceId );
         for( auto&& it : set )
            ( *it ).synchronize();
         return true;
      }
      catch( const std::out_of_range& ) {
         return false;
      }
   }

protected:
   using SetType = std::unordered_set< SmartPointer* >;

   std::unordered_map< int, SetType > pointersOnDevices;
};

// TODO: Device -> Allocator (in all smart pointers)
template< typename Device >
SmartPointersRegister&
getSmartPointersRegister()
{
   static SmartPointersRegister reg;
   return reg;
}

template< typename Device >
Timer&
getSmartPointersSynchronizationTimer()
{
   static Timer timer;
   return timer;
}

/**
 * Negative deviceId means that the ID of the currently active device will be
 * determined automatically.
 */
template< typename Device >
bool
synchronizeSmartPointersOnDevice( int deviceId = -1 )
{
   // TODO: better way to skip synchronization of host-only smart pointers
   if constexpr( std::is_same_v< Device, Devices::Sequential > || std::is_same_v< Device, Devices::Host > ) {
      return true;
   }
   else {
      getSmartPointersSynchronizationTimer< Device >().start();
      bool b = getSmartPointersRegister< Device >().synchronizeDevice( deviceId );
      getSmartPointersSynchronizationTimer< Device >().stop();
      return b;
   }
}

}  // namespace TNL::Pointers
