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

#ifndef PUB_SUB2_PUB_TCP_HH
#define PUB_SUB2_PUB_TCP_HH

#include "pub_intl.hh"

#include <future>
#include <iostream>
#include <thread>

namespace pub_sub
{
    namespace detail
    {
        /*!
         * @brief PubClientTcp tracks client connections (subscribers) to a
         * Publisher
         */
        class PubClientTcp
        {
        public:
            /*!
             * @brief create a PubClientTcp
             * @param context The io_context to work with
             * @param messages The history of recent messages
             * @param sleep_list A place to register the client when there are
             * no messages to send
             * @param debug Debug/stats hooks
             */
            explicit PubClientTcp(
                boost::asio::io_context&                             context,
                MessageHistory&                                      messages,
                std::vector< boost::intrusive_ptr< PubClientTcp > >& sleep_list,
                PubDebugNotices&                                     debug )
                : s_{ context }, deadline_( context ), messages_{ messages },
                  sleep_list_{ sleep_list }, debug_{ debug },
                  initial_req_{ NEXT_PUB_MSG( ) },
                  initial_req_buffer_( boost::asio::mutable_buffer(
                      &initial_req_, sizeof( initial_req_ ) ) )
            {
                deadline_.expires_at(
                    boost::asio::steady_timer::time_point::max( ) );
            }
            ~PubClientTcp( )
            {
                std::cout << "Closing client" << std::endl;
            }

            tcp::socket&
            socket( ) noexcept
            {
                return s_;
            }
            int
            ref_count( ) const noexcept
            {
                return ref_cnt_;
            }

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

            /*!
             * @brief Start the client handler
             * @details The system works in the following scheme
             * 1. read the request from the subscriber
             * 2. determine which message to start sending
             * 3. send messages until the client has disconnected
             */
            void
            run( )
            {
                std::cout << "Client running refcount " << ref_count( )
                          << std::endl;

                tcp::no_delay             option( true );
                boost::system::error_code ignored{ };
                s_.set_option( option, ignored );

                set_timeout( );
                deadline_.async_wait(
                    [ self = shared_from_this( ) ](
                        const boost::system::error_code& ec ) {
                        self->timer_handler( ec );
                    } );
                boost::asio::async_read(
                    s_,
                    initial_req_buffer_,
                    [ self = shared_from_this( ) ](
                        const boost::system::error_code& ec,
                        std::size_t                      bytes_transferred ) {
                        if ( !ec )
                        {
                            self->clear_timeout( );
                            self->dispatch_initial_request( );
                        }
                    } );
            }

            /*!
             * @brief schedule an async write if possible
             * @details A client that has fallen behind the message history
             * is skipped ahead to the oldest retained message.
             * @return True if a write could be scheduled, else false
             */
            bool
            schedule_write( )
            {
                if ( !has_position_ )
                {
                    auto cur = messages_.newest( );
                    if ( !cur )
                    {
                        return false;
                    }
                    write_message( std::move( cur ) );
                    return true;
                }
                auto cur = messages_.first_after( last_key_ );
                if ( !cur )
                {
                    return false;
                }
                if ( messages_.has_gap_after( last_key_ ) )
                {
                    debug_.slow_client_detected( last_key_, cur->key( ) );
                }
                write_message( std::move( cur ) );
                return true;
            }

            /*!
             * @brief schedule a async write if possible, else put the client
             * in the sleep list.
             */
            void
            schedule_write_or_sleep( )
            {
                if ( schedule_write( ) )
                {
                    return;
                }
                sleep( );
            }
            /*!
             * @brief put the client on the sleep list
             */
            void
            sleep( )
            {
                /* Find the first empty entry and use it it possible */
                auto it = std::find_if(
                    sleep_list_.begin( ),
                    sleep_list_.end( ),
                    []( const boost::intrusive_ptr< PubClientTcp >& cur )
                        -> bool { return cur.get( ) == nullptr; } );
                if ( it == sleep_list_.end( ) )
                {
                    sleep_list_.emplace_back( shared_from_this( ) );
                }
                else
                {
                    *it = shared_from_this( );
                }
            }

        private:
            void
            clear_timeout( )
            {
                deadline_.expires_at(
                    boost::asio::steady_timer::time_point::max( ) );
            }
            void
            set_timeout( )
            {
                deadline_.expires_after( std::chrono::seconds( 2 ) );
            }

            void
            timer_handler( const boost::system::error_code& ec )
            {
                if ( deadline_.expiry( ) <=
                     boost::asio::steady_timer::clock_type::now( ) )
                {
                    std::cout << "Client timed out, closing the socket"
                              << std::endl;
                    boost::system::error_code ignored{ };
                    s_.cancel( ignored );
                    s_.close( ignored );
                }
                else
                {
                    deadline_.async_wait(
                        [ self = shared_from_this( ) ](
                            const boost::system::error_code& ec ) {
                            self->timer_handler( ec );
                        } );
                }
            }

            /*!
             * @brief startup handler/dispatcher for the client.
             * This determines which message to start with.
             */
            void
            dispatch_initial_request( )
            {
                if ( initial_req_.request == NEXT_PUB_MSG( ) )
                {
                    has_position_ = false;
                    sleep( );
                    return;
                }
                /* a request for a key that predates the retained history
                 * is skipped ahead by schedule_write like any other gap */
                last_key_ = initial_req_.request;
                has_position_ = true;
                schedule_write_or_sleep( );
            }

            /*!
             * @brief Start the async write operation for the message
             * @param cur The message to write
             */
            void
            write_message( MessageBlockPtr cur )
            {
                if ( !cur )
                {
                    return;
                }
                in_flight_ = std::move( cur );
                last_key_ = in_flight_->key( );
                has_position_ = true;
                set_timeout( );
                boost::asio::async_write(
                    s_,
                    in_flight_->buffers( ),
                    [ self2 = shared_from_this( ) ](
                        const boost::system::error_code& ec,
                        std::size_t                      bytes_transferred ) {
                        /* release the payload before looking for more work */
                        self2->in_flight_.reset( );
                        if ( ec )
                        {
                            return;
                        }
                        self2->clear_timeout( );
                        self2->schedule_write_or_sleep( );
                    } );
            }

            int                       ref_cnt_{ 0 };
            tcp::socket               s_;
            boost::asio::steady_timer deadline_;
            KeyType                   last_key_{ 0 };
            bool                      has_position_{ false };
            MessageBlockPtr           in_flight_{ nullptr };

            MessageHistory&                                      messages_;
            std::vector< boost::intrusive_ptr< PubClientTcp > >& sleep_list_;
            PubDebugNotices&                                     debug_;

            SubcriptionRequest          initial_req_;
            boost::asio::mutable_buffer initial_req_buffer_;

            friend void intrusive_ptr_add_ref( PubClientTcp* );
            friend void intrusive_ptr_release( PubClientTcp* );
        };

        using PubClientPtr = boost::intrusive_ptr< PubClientTcp >;
        inline PubClientPtr
        make_pub_client( boost::asio::io_context&     context,
                         MessageHistory&              messages,
                         std::vector< PubClientPtr >& sleep_list,
                         PubDebugNotices&             debug )
        {
            auto p = std::make_unique< PubClientTcp >(
                context, messages, sleep_list, debug );
            return p.release( );
        }

        inline void
        intrusive_ptr_add_ref( PubClientTcp* pc )
        {
            pc->ref_cnt_++;
        }

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

        /*!
         * @brief The PublisherIntlTcp is the logic behind the Publisher.
         * @details This is used as a compilation firewall to keep the
         * users of pub_sub from directly seeing the underlying I/O interface
         * (ie asio).
         */
        struct PublisherIntlTcp : public plugins::PublisherInstance
        {
            PublisherIntlTcp( const std::string& address,
                              int                port,
                              PubDebugNotices&   debug,
                              PublisherOptions   options =
                                  PublisherOptions::from_environment( ) )
                : context_{ },
                  acceptor_{ context_,
                             tcp::endpoint(
                                 boost::asio::ip::make_address( address ),
                                 port ) },
                  stopping_{ false }, messages_{ options.max_depth },
                  debug_{ debug }
            {
                sleeping_clients_.reserve( 100 );
                context_.post( [ this ]( ) { accept_loop( ); } );
                thread_ = std::thread( [ this ]( ) { context_.run( ); } );
            }

            ~PublisherIntlTcp( ) override
            {
                context_.post( [ this ]( ) {
                    stopping_ = true;
                    wake_up_clients( );
                    context_.stop( );
                } );
                thread_.join( );
            }

            /*!
             * @brief Start of the accept sequence
             */
            void
            accept_loop( )
            {
                if ( stopping_ )
                {
                    return;
                }
                auto client = make_pub_client(
                    context_, messages_, sleeping_clients_, debug_ );
                acceptor_.async_accept(
                    client->socket( ),
                    [ client,
                      this ]( const boost::system::error_code& ec ) mutable {
                        handle_accept( std::move( client ), ec );
                    } );
            }

            /*!
             * @brief End of the accept sequence, queues another accept on
             * completion
             * @param pc The new client
             * @param ec Error code
             */
            void
            handle_accept( PubClientPtr                     pc,
                           const boost::system::error_code& ec )
            {
                if ( stopping_ )
                {
                    return;
                }
                if ( !ec )
                {
                    pc->run( );
                }
                accept_loop( );
            }

            /*!
             * @brief wake up all sleeping clients, call when new data has
             * arrived
             */
            void
            wake_up_clients( )
            {
                if ( stopping_ )
                {
                    sleeping_clients_.clear( );
                    return;
                }
                std::for_each( sleeping_clients_.begin( ),
                               sleeping_clients_.end( ),
                               []( PubClientPtr& p ) {
                                   if ( p.get( ) == nullptr )
                                   {
                                       return;
                                   }
                                   if ( p->schedule_write( ) )
                                   {
                                       p = nullptr;
                                   }
                               } );
                auto it = std::find_if( sleeping_clients_.begin( ),
                                        sleeping_clients_.end( ),
                                        []( const PubClientPtr& cur ) -> bool {
                                            return cur.get( ) != nullptr;
                                        } );
                if ( it == sleeping_clients_.end( ) )
                {
                    sleeping_clients_.clear( );
                }
            }

            tcp::endpoint
            local_endpoint( )
            {
                auto promise =
                    std::make_shared< std::promise< tcp::endpoint > >( );
                auto result = promise->get_future( );
                context_.post( [ this, promise ]( ) {
                    promise->set_value( acceptor_.local_endpoint( ) );
                } );
                return result.get( );
            }

            /*!
             * @brief insert a message into the publishing queue
             * @param key Message key
             * @param msg Message buffer
             * @note this is called on the background I/O thread
             */
            void
            insert_message( KeyType key, Message msg )
            {
                messages_.push( key, std::move( msg ) );
                wake_up_clients( );
            }

            /*!
             * @brief Insert a message to be published
             * @param key
             * @param msg
             * @note This is called from the main application thread and
             * sends its data in a thread safe manner to the background thread
             */
            void
            publish( KeyType key, Message msg ) override
            {
                if ( msg.length > MAX_MSG_SIZE( ) )
                {
                    throw std::invalid_argument( "Message too big to send" );
                }
                context_.post( [ this, key, msg ]( ) mutable {
                    insert_message( key, std::move( msg ) );
                } );
            }
            boost::asio::io_context     context_;
            tcp::acceptor               acceptor_;
            std::thread                 thread_;
            bool                        stopping_;
            MessageHistory              messages_;
            std::vector< PubClientPtr > sleeping_clients_{ };
            PubDebugNotices&            debug_;
        };
    } // namespace detail
} // namespace pub_sub

#endif
