/** \file
 *
 *  Contains utilities to aid ExternalSimGazebo class implementation.
 *
 *  Copyright (c) 2022 MBARI
 *  MBARI Proprietary Information.  All Rights Reserved
 */
#ifndef EXTERNALSIMGAZEBOUTILS_H_
#define EXTERNALSIMGAZEBOUTILS_H_

#include <chrono>
#include <deque>
#include <functional>
#include <iostream>
#include <memory>
#include <mutex>
#include <string>
#include <utility>
#include <variant>

#include <gz/msgs/Utility.hh>
#include <gz/transport/Node.hh>

/**
 *  Subscribes to more than one Gazebo Transport topic
 *  simultaneously and synchronizes messages for ease of
 *  use. It is assumed that messages have a header and
 *  that perfect synchronization is achievable (e.g. when
 *  published from within Gazebo Gazebo's event loop).
 *
 *  ExternalSimGazeboUtils.h should only be included by
 *  ExternalSimGazebo.h and/or ExternalSimGazebo.cpp
 *
 *  \ingroup modules_simulator
 */
class SynchronizedGazeboSubscriber
{
public:
    /// Subscribe to `topic0` and `topic1`, registering a `callback`
    /// \param[in] node Gazebo Transport node to use for subscription
    /// \param[in] topic0 First topic name
    /// \param[in] topic1 Second topic name
    /// \param[in] maxBufferSize Cap to buffer size when synchronizing
    /// \param[in] callback Callable with the following parameters:
    ///   \param[in] msg0 First topic message
    ///   \param[in] msg1 Second topic message
    /// \return a SynchronizedGazeboSubscriber instance
    template<typename MessageT0, typename MessageT1>
    static SynchronizedGazeboSubscriber Subscribe(
        gz::transport::Node& node,
        const std::string& topic0, const std::string& topic1, size_t maxBufferSize,
        std::function<void( const MessageT0& msg0, const MessageT1& msg1 )> callback )
    {
        using ImplementationT = ImplementationForTwoTopics<MessageT0, MessageT1>;
        std::unique_ptr<ImplementationT> impl =
            std::make_unique<ImplementationT>( std::move( callback ), maxBufferSize );
        auto callback0 = &ImplementationT::OnMessage0;
        if( !node.Subscribe( topic0, callback0, impl.get() ) )
        {
            return SynchronizedGazeboSubscriber();
        }
        auto callback1 = &ImplementationT::OnMessage1;
        if( !node.Subscribe( topic1, callback1, impl.get() ) )
        {
            if( !node.Unsubscribe( topic0 ) )
            {
                std::cerr << "Failed to unsubscribe from " << topic0
                          << " after failing to subscribe to " << topic1
                          << std::endl;
            }
            return SynchronizedGazeboSubscriber();
        }
        return SynchronizedGazeboSubscriber( std::move( impl ) );
    }

    /// Subscribe to `topic0` and `topic1`, registering a `callback`
    /// \param[in] node Gazebo Transport node to use for subscription
    /// \param[in] topic0 First topic name
    /// \param[in] topic1 Second topic name
    /// \param[in] maxBufferSize Cap to buffer size when synchronizing
    /// \param[in] callback Pointer to class instance method with the
    ///   following parameters:
    ///     \param[in] msg0 First topic message
    ///     \param[in] msg1 Second topic message
    /// \param[in] obj Pointer to class instance
    /// \return a SynchronizedGazeboSubscriber instance
    template<typename ClassT, typename MessageT0, typename MessageT1>
    static SynchronizedGazeboSubscriber Subscribe(
        gz::transport::Node& node,
        const std::string& topic0, const std::string& topic1, size_t maxBufferSize,
        void( ClassT::*callback )( const MessageT0& msg0, const MessageT1& msg1 ),
        ClassT *obj )
    {
        using namespace std::placeholders;
        return Subscribe<MessageT0, MessageT1>(
                   node, topic0, topic1, maxBufferSize,
                   std::bind( callback, obj, _1, _2 ) );
    }

    /// Default constructor
    SynchronizedGazeboSubscriber() = default;

    /// Check whether this subscriber is valid
    /// i.e. active and synchronizing messages
    /// \return true if valid, false otherwise
    bool Valid() const
    {
        return impl_ != nullptr;
    }

    /// Bool operator overload
    /// \see SynchronizedGazeboSubscriber::Valid()
    explicit operator bool() const
    {
        return Valid();
    }

private:
    /// Implementation interface
    class Implementation
    {
    public:
        virtual ~Implementation() = default;
    };

    /// Implementation for synchronizing two topics
    template<typename MessageT0, typename MessageT1>
    class ImplementationForTwoTopics : public Implementation
    {
    public:
        using CallbackT =
            std::function<void( const MessageT0&, const MessageT1& )>;

        /// Constructor
        /// \param[in] callback Synchronized subscriber callback
        /// \param[in] maxBufferSize Buffer size cap
        ImplementationForTwoTopics( CallbackT callback, size_t maxBufferSize )
            : callback_( std::move( callback ) ), maxBufferSize_( maxBufferSize )
        {
        }

        /// Callback to notify arrival of a message of the first type
        void OnMessage0( const MessageT0& message0 )
        {
            MessageT1 message1;
            if( Synchronize( message0, &message1 ) )
            {
                callback_( message0, message1 );
            }
        }

        /// Callback to notify arrival of a message of the second type
        void OnMessage1( const MessageT1& message1 )
        {
            MessageT0 message0;
            if( Synchronize( message1, &message0 ) )
            {
                callback_( message0, message1 );
            }
        }

    private:
        /// Synchronize `message0`, checking if a synchronization
        /// point has been reached i.e. if there's a message of the
        /// opposite type with a matchin stamp in the buffer.
        ///
        /// \param[in] message0 message to synchronize
        /// \param[out] message1 to be populated when a
        ///   synchronization point is reached.
        /// \return true if a synchronization point has
        ///   been reached, false otherwise
        template<typename MessageU0, typename MessageU1>
        bool Synchronize( const MessageU0& message0, MessageU1* message1 )
        {
            static_assert(
                ( std::is_same_v<MessageU0, MessageT0> &&
                  std::is_same_v<MessageU1, MessageT1> ) ||
                ( std::is_same_v<MessageU1, MessageT0> &&
                  std::is_same_v<MessageU0, MessageT1> ) );
            const std::chrono::steady_clock::duration stamp =
                gz::msgs::Convert( message0.header().stamp() );
            std::lock_guard<std::mutex> lock( bufferMutex_ );
            // Find event stamped at a time not earlier
            // than the current stamp.
            auto it = std::lower_bound(
                          buffer_.begin(), buffer_.end(), stamp,
                          []( const DelayedReceptionEvent & event,
                              const std::chrono::steady_clock::duration & stamp )
            {
                return event.stamp < stamp;
            } );
            if( it == buffer_.end() )
            {
                // All buffered events are stamped at an earlier time.
                // First time we get a stamp this much into the future.
                // Buffer message for delayed reception.
                buffer_.push_back( {stamp, message0} );
                if( buffer_.size() > maxBufferSize_ )
                {
                    buffer_.pop_front();
                }
                return false;
            }
            if( it->stamp > stamp )
            {
                // No event found with a matching stamp.
                // Publishers out of sync? Ignore message
                // and drop all earlier events.
                buffer_.erase( buffer_.begin(), it );
                return false;
            }
            if( std::holds_alternative<MessageU0>( it->message ) )
            {
                std::cerr << "Duplicate timestamps found while "
                          << "synchronizing messages. Dropping "
                          << "earlier messages." << std::endl;
                it->message = message0;
                return false;
            }
            // Got an event with a matching stamp.
            // Move it out of the buffer and drop
            // all earlier events before notifying.
            *message1 = std::get<MessageU1>( std::move( it->message ) );
            buffer_.erase( buffer_.begin(), ++it );
            return true;
        }

        /// Synchronized subscriber callback
        CallbackT callback_;

        /// A reception event of a message of either type
        struct DelayedReceptionEvent
        {
            std::chrono::steady_clock::duration stamp;
            std::variant<MessageT0, MessageT1> message;
        };

        /// Buffer for message synchronization
        std::deque<DelayedReceptionEvent> buffer_;

        /// Buffer size cap
        const size_t maxBufferSize_;

        /// Mutex for thread-safe buffer access
        std::mutex bufferMutex_;
    };

    /// Constructor
    /// \param[in] impl Type-erased implementation to be used
    SynchronizedGazeboSubscriber( std::unique_ptr<Implementation> impl )
        : impl_( std::move( impl ) )
    {
    }

    /// Private, type-erased implementation
    std::unique_ptr<Implementation> impl_{nullptr};
};

#endif /* EXTERNALSIMGAZEBOUTILS_H_ */