//
// Created by jonathan.hanks on 2/20/20.
//

#ifndef PUB_SUB2_SUB_TCP_HH
#define PUB_SUB2_SUB_TCP_HH

#include <atomic>
#include <functional>
#include <iostream>
#include <memory>

#include <cds-pubsub/sub_plugin.hh>
#include <boost/smart_ptr/intrusive_ptr.hpp>
#include "sub_intl.hh"

namespace pub_sub
{
    namespace detail
    {
        /*!
         * @brief the ClientTermNotifier is used to notify Subscriptions that
         * an individual connection has closed (giving it the last message
         * recieved).
         */
        using ClientTermNotifier = std::function< void( KeyType ) >;

        /*!
         * @brief SubClients represent individual connections to the publisher
         */
        class SubClient
        {
        public:
            /*!
             * @brief Create a connection to a publisher
             * @note Do not directly call, this must be created through
             * make_sub_client
             * @param context The io_context to work with
             * @param last_msg The last message received
             * @param endpoint The endpoint to connect to
             * @param sub_id SubscriptionTcp id
             * @param handler SubscriptionTcp handler
             * @param on_terminate Callback to notify the higher level
             * subscription that we are exited
             */
            SubClient( io_context&          context,
                       KeyType              last_msg,
                       tcp::endpoint&       endpoint,
                       SubId                sub_id,
                       SubHandler&          handler,
                       SubDebugNotices&     debug,
                       ClientTermNotifier   on_terminate,
                       std::atomic< bool >& stopping )
                : s_{ context }, last_msg_{ last_msg }, sub_id_{ sub_id },
                  handler_{ handler }, debug_{ debug },
                  on_terminate_{ std::move( on_terminate ) },
                  stopping_{ stopping }, sub_req_{ last_msg_ },
                  header_{ NEXT_PUB_MSG( ), 0 }, timer_{ context }
            {
                //                std::cout << "Create sub client addr: " <<
                //                (void*)this
                //                          << std::endl;

                s_.async_connect( endpoint,
                                  [ self = shared_from_this( ) ](
                                      const boost::system::error_code& ec ) {
                                      self->timer_.cancel( );
                                      if ( ec )
                                      {
                                          return;
                                      }
                                      self->connected_ = true;
                                      self->request_message( );
                                  } );
                start_timeout_timer();
            }

            ~SubClient( )
            {
                timer_.cancel( );
                //                std::cout << "sub client terminating last key
                //                = " << last_msg_
                //                          << " addr: " << (void*)this <<
                //                          std::endl;
                if ( !stopping_ )
                {
                    on_terminate_( last_msg_ );
                }
            }

            boost::intrusive_ptr< SubClient >
            shared_from_this( )
            {
                return this;
            }

            boost::asio::ip::tcp::socket*
            get_socket( )
            {
                return &s_;
            }

        private:
            void
            start_timeout_timer()
            {
                timer_.expires_from_now( timeout_ );
                timer_.async_wait( [self = shared_from_this( )](
                        const boost::system::error_code& ec ) {
                    self->connection_timeout( ec );
                } );
            }
            void
            connection_timeout( const boost::system::error_code& ec )
            {
                if (ec)
                {
                    return;
                }
                boost::system::error_code ignored{ };
                s_.close( ignored );
            }
            /*!
             * @brief First step in the subscription is to request data
             */
            void
            request_message( )
            {
                boost::system::error_code ignored{ };
                {
                    tcp::no_delay option( true );
                    s_.set_option( option, ignored );
                }
                {
                    boost::asio::socket_base::receive_buffer_size option;
                    s_.get_option( option );
                    std::cout
                        << "Socket recieve buffer size is: " << option.value( )
                        << std::endl;
                    if ( option.value( ) < 2 * 1024 * 1024 )
                    {
                        option = 2 * 1024 * 1024;
                        s_.set_option( option, ignored );
                        s_.get_option( option );
                        std::cout << "Socket recieve buffer size is: "
                                  << option.value( ) << std::endl;
                    }
                }
                debug_.connection_started( sub_id_, last_msg_ );
                sub_req_.request = last_msg_;
                boost::asio::const_buffer buf( &sub_req_, sizeof( sub_req_ ) );
                boost::asio::async_write(
                    s_,
                    buf,
                    [ self = shared_from_this( ) ](
                        const boost::system::error_code& ec,
                        std::size_t                      bytes_transferred ) {
                        self->timer_.cancel();
                        if ( ec )
                        {
                            return;
                        }
                        self->read_message_header( );
                    } );
                start_timeout_timer();
            }

            /*!
             * @brief the start of the read loop, read the headers
             */
            void
            read_message_header( )
            {
                boost::asio::mutable_buffer buf( &header_, sizeof( header_ ) );
                boost::asio::async_read(
                    s_,
                    buf,
                    [ self = shared_from_this( ) ](
                        const ::boost::system::error_code& ec,
                        std::size_t                        bytes_transferred ) {
                        self->timer_.cancel();
                        if ( ec )
                        {
                            return;
                        }
                        self->read_message_body( );
                    } );
                start_timeout_timer();
            }

            /*!
             * @brief After reading the headers prepare the message buffers and
             * request the message body
             */
            void
            read_message_body( )
            {
                if ( header_.size > MAX_MSG_SIZE( ) )
                {
                    return;
                }

                TimePoint start = std::chrono::steady_clock::now( );
                DataPtr   data( new unsigned char[ header_.size ] );
                boost::asio::mutable_buffer buf( data.get( ), header_.size );
                boost::asio::async_read(
                    s_,
                    buf,
                    [self = shared_from_this( ),
                     data_buf = std::move( data ),
                     start]( const ::boost::system::error_code& ec,
                             std::size_t bytes_transferred ) mutable {
                        self->timer_.cancel();
                        if ( ec )
                        {
                            return;
                        }
                        self->receive_message_body( std::move( data_buf ),
                                                    start );
                    } );
                start_timeout_timer();
            }

            /*!
             * @brief Handle a received message body then restart the read loop
             * @param data The message body.
             */
            void
            receive_message_body( DataPtr data, TimePoint start_time )
            {
                SubMessage sub_msg(
                    sub_id_,
                    header_.key,
                    Message( std::move( data ), header_.size ) );
                // std::cout << "recieved message " << header_.key << " ("
                //          << header_.size << ")" << std::endl;
                auto duration =
                    std::chrono::duration_cast< std::chrono::milliseconds >(
                        std::chrono::steady_clock::now( ) - start_time );
                debug_.message_received(
                    sub_id_, header_.key, header_.size, duration );

                handler_( std::move( sub_msg ) );
                last_msg_ = header_.key;
                read_message_header( );
            }

            int                       ref_cnt_{ 0 };
            tcp::socket               s_;
            KeyType                   last_msg_;
            SubId                     sub_id_;
            SubHandler&               handler_;
            SubDebugNotices&          debug_;
            ClientTermNotifier        on_terminate_;
            SubcriptionRequest        sub_req_;
            BlockHeader               header_;
            std::atomic< bool >&      stopping_;
            bool                      connected_{ false };
            boost::asio::steady_timer timer_;
            std::chrono::duration<std::int64_t> timeout_{ std::chrono::seconds(60) };

            friend void intrusive_ptr_add_ref( SubClient* );
            friend void intrusive_ptr_release( SubClient* );
        };
        using SubClientPtr = boost::intrusive_ptr< SubClient >;
        inline SubClientPtr
        make_sub_client( io_context&          context,
                         KeyType              last_msg,
                         tcp::endpoint&       endpoint,
                         SubId                sub_id,
                         SubHandler&          handler,
                         SubDebugNotices&     debug,
                         ClientTermNotifier   on_terminate,
                         std::atomic< bool >& stopping )
        {
            auto s = std::make_unique< SubClient >( context,
                                                    last_msg,
                                                    endpoint,
                                                    sub_id,
                                                    handler,
                                                    debug,
                                                    std::move( on_terminate ),
                                                    stopping );
            return s.release( );
        }
        inline void
        intrusive_ptr_add_ref( SubClient* sc )
        {
            sc->ref_cnt_++;
        }

        inline void
        intrusive_ptr_release( SubClient* sc )
        {
            sc->ref_cnt_--;
            if ( sc->ref_cnt_ <= 0 )
            {
                std::unique_ptr< SubClient > d{ sc };
            }
        }

        /*!
         * @brief High level SubscriptionTcp handler
         * @details This manages the startup/retry of the connections
         */
        class SubscriptionTcp : public pub_sub::plugins::Subscription
        {
        public:
            SubscriptionTcp( ) = delete;

            SubscriptionTcp( io_context&          context,
                             tcp::endpoint        endpoint,
                             SubHandler           handler,
                             std::atomic< bool >& stopping,
                             SubDebugNotices&     debug )
                : Subscription( ), context_{ context }, endpoint_{ std::move(
                                                            endpoint ) },
                  handler_{ std::move( handler ) }, debug_{ debug },
                  last_key_{ NEXT_PUB_MSG( ) }, stopping_{ stopping },
                  wait_state_{ pauses( ).cbegin( ) }, timer_{ context }
            {
                std::cout << "SubscriptionTcp against " << endpoint_ << " : "
                          << (void*)this << std::endl;
                do_launch_client( );
            }

            ~SubscriptionTcp( ) override
            {
                std::cout << "SubscriptionTcp destroyed :" << (void*)this
                          << std::endl;
            }

            SubscriptionTcp( const SubscriptionTcp& ) = delete;
            SubscriptionTcp( SubscriptionTcp&& ) = delete;
            SubscriptionTcp& operator=( const SubscriptionTcp& ) = delete;
            SubscriptionTcp& operator=( SubscriptionTcp&& ) = delete;

            void
            reset( ) override
            {
                context_.post( [ this ]( ) {
                    if ( s_ )
                    {
                        s_->cancel( );
                        wait_state_ = pauses( ).begin( );
                        last_key_ = NEXT_PUB_MSG( );
                    }
                } );
            }

        private:
            /*!
             * @brief determine how long to wait before trying to reconnect
             * @param key The last key returned by the previous connection
             * or NEXT_PUB_MSG()
             * @return The duration of time to wait before attempting to
             * reconnect
             */
            std::chrono::milliseconds
            next_wait( KeyType key )
            {
                auto wait = std::chrono::milliseconds( 0 );
                /* If key != last_key_ we made progress, so reset
                 * the waits to small values
                 */
                if ( key != last_key_ )
                {
                    last_key_ = key;
                    wait_state_ = pauses( ).begin( );
                    return wait;
                }
                if ( wait_state_ != pauses( ).end( ) )
                {
                    wait = *wait_state_;
                    ++wait_state_;
                }
                else
                {
                    wait = pauses( ).back( );
                }
                return wait;
            }
            /*!
             * @brief launch a client, doing a wait if needed
             * @param key the key of the last message processed
             * or NEXT_PUB_MSG()
             */
            void
            launch_new_client( KeyType key )
            {
                if ( stopping_ )
                {
                    return;
                }
                auto wait = next_wait( key );
                if ( wait.count( ) == 0 )
                {
                    do_launch_client( );
                    return;
                }
                std::cerr << "waiting " << wait.count( ) << "ms\n";
                timer_.expires_after( wait );
                timer_.async_wait(
                    [ this ]( const boost::system::error_code& ec ) {
                        do_launch_client( );
                    } );
            }

            /*!
             * @brief The does the actual launching of a client
             */
            void
            do_launch_client( )
            {
                if ( stopping_ )
                {
                    return;
                }
                //                std::cout << "this is = " << (void*)( this )
                //                << std::endl;
                auto id = id_;
                auto sub_client = make_sub_client(
                    context_,
                    last_key_,
                    endpoint_,
                    id,
                    handler_,
                    debug_,
                    [ this ]( KeyType new_key ) {
                        s_ = nullptr;
                        if ( stopping_ )
                        {
                            return;
                        }
                        debug_.connection_terminated( sub_id( ), new_key );
                        // std::cout
                        //    << "on_terminate called with this = " <<
                        //    (void*)this
                        //    << " new_key = " << new_key << std::endl;
                        this->launch_new_client( new_key );
                    },
                    stopping_ );
                s_ = sub_client->get_socket( );
            }

            /*!
             * @brief A set of pauses to allow a for problems/congestion/... to
             * clear between reconnection attempts.
             * @return The set of pauses
             */
            const static std::vector< std::chrono::milliseconds >&
            pauses( )
            {
                const static std::vector< std::chrono::milliseconds > p{
                    std::chrono::milliseconds( 0 ),
                    std::chrono::milliseconds( 5 ),
                    std::chrono::milliseconds( 50 ),
                    std::chrono::milliseconds( 500 ),
                    std::chrono::seconds( 1 ),
                    std::chrono::seconds( 5 ),
                    std::chrono::seconds( 15 ),
                    // std::chrono::seconds(30),
                };
                return p;
            }

            io_context&                   context_;
            tcp::endpoint                 endpoint_;
            SubHandler                    handler_;
            KeyType                       last_key_;
            std::atomic< bool >&          stopping_;
            boost::asio::ip::tcp::socket* s_{ nullptr };
            SubDebugNotices&              debug_;
            std::vector< std::chrono::milliseconds >::const_iterator
                                      wait_state_;
            boost::asio::steady_timer timer_;
        };
    } // namespace detail
} // namespace pub_sub

#endif // PUB_SUB2_SUB_TCP_HH
