//
// 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 list of messages
             * @param sleep_list A place to register the client when there are
             * no messages to send
             */
            explicit PubClientTcp(
                boost::asio::io_context&     context,
                BoundedList< MessageBlock >& messages,
                std::vector< boost::intrusive_ptr< PubClientTcp > >&
                    sleep_list )
                : s_{ context }, deadline_( context ),
                  last_message_{ nullptr }, messages_{ messages },
                  sleep_list_{ sleep_list }, 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 )
                        {
                            //                            std::cout << "socket "
                            //                                      <<
                            //                                      self->socket(
                            //                                      ).native_handle(
                            //                                      )
                            //                                      << "
                            //                                      refcount "
                            //                                      <<
                            //                                      self->ref_count(
                            //                                      )
                            //                                      << " run
                            //                                      lambda" <<
                            //                                      std::endl;
                            self->clear_timeout( );
                            self->dispatch_initial_request( );
                        }
                    } );
            }

            /*!
             * @brief schedule an async write if possible
             * @return True if a write could be scheduled, else false
             */
            bool
            schedule_write( )
            {
                //                std::cout << "socket " << socket(
                //                ).native_handle( )
                //                          << " sched write lm = "
                //                          << ( last_message_ ? "exists" :
                //                          "null" ) << std::endl;
                MessageBlockPtr cur{ ( last_message_ ? last_message_->next( )
                                                     : nullptr ) };
                //                std::cout << "cur is " << ( cur ? "not null" :
                //                "null" )
                //                          << " refcount " << ref_count( ) << "
                //                          schedule write"
                //                          << std::endl;
                if ( !cur )
                {
                    cur = messages_.head( );
                    if ( !cur )
                    {
                        return false;
                    }
                    if ( last_message_ && last_message_->key( ) <= cur->key( ) )
                    {
                        return false;
                    }
                }
                write_message( cur );
                return true;
            }

            /*!
             * @brief schedule a async write if possible, else put the client
             * in the sleep list.
             */
            void
            schedule_write_or_sleep( )
            {
                //                std::cout << "socket " << socket(
                //                ).native_handle( )
                //                          << " last_msg_ = " <<
                //                          this->last_message_->key( )
                //                          << " refcount " << ref_count( )
                //                          << " schedule_write_or_sleep" <<
                //                          std::endl;
                if ( schedule_write( ) )
                {
                    return;
                }
                sleep( );
            }
            /*!
             * @brief put the client on the sleep list
             */
            void
            sleep( )
            {
                //                std::cout << "socket " << socket(
                //                ).native_handle( )
                //                          << " refcount " << ref_count( ) << "
                //                          sleep pre"
                //                          << std::endl;
                /* 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( );
                }
                //                std::cout << "socket " << socket(
                //                ).native_handle( )
                //                          << " refcount " << ref_count( ) << "
                //                          sleep post"
                //                          << std::endl;
            }

        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;
                    s_.cancel( );
                    s_.close( );
                }
                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 ( messages_.empty( ) ||
                     initial_req_.request == NEXT_PUB_MSG( ) )
                {
                    sleep( );
                    return;
                }
                auto cur = messages_.tail( ).get( );
                while ( cur )
                {
                    if ( cur->key( ) > initial_req_.request )
                    {
                        write_message( cur );
                        return;
                    }
                    cur = cur->next( ).get( );
                }
                sleep( );
            }

            /*!
             * @brief Start the async write operation for the message
             * @param cur The message to write
             */
            void
            write_message( MessageBlockPtr cur )
            {
                last_message_ = std::move( cur );
                if ( !last_message_ )
                {
                    return;
                }
                //                std::cout << "Sent message " <<
                //                last_message_->key( )
                //                          << " to socket " <<
                //                          s_.native_handle( )
                //                          << " refcount = " << ref_count( ) <<
                //                          std::endl;

                //                s_.get_io_context().post([self=shared_from_this()](){
                //                    std::cout << "Refcount = " <<
                //                    self->ref_count() << " write_message outer
                //                    lambda" << std::endl;
                set_timeout( );
                boost::asio::async_write(
                    s_,
                    last_message_->asio_buffers( ),
                    [self2 = shared_from_this( )](
                        const boost::system::error_code& ec,
                        std::size_t                      bytes_transferred ) {
                        if ( ec )
                        {
                            //                            std::cout << "Error on
                            //                            write" << std::endl;
                            //                            std::cout << "Refcount
                            //                            = " <<
                            //                            self2->ref_count( )
                            //                                      << "
                            //                                      write_message
                            //                                      lambda
                            //                                      error"
                            //                                      <<
                            //                                      std::endl;
                            return;
                        }
                        //                        std::cout << "socket "
                        //                                  << self2->socket(
                        //                                  ).native_handle( )
                        //                                  << " refcount " <<
                        //                                  self2->ref_count( )
                        //                                  << " write_message
                        //                                  lambda" <<
                        //                                  std::endl;
                        self2->clear_timeout( );
                        self2->schedule_write_or_sleep( );
                    } );
                //                });
            }

            int                       ref_cnt_{ 0 };
            tcp::socket               s_;
            boost::asio::steady_timer deadline_;
            MessageBlockPtr           last_message_;

            BoundedList< MessageBlock >&                         messages_;
            std::vector< boost::intrusive_ptr< PubClientTcp > >& sleep_list_;

            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,
                         BoundedList< MessageBlock >& messages,
                         std::vector< PubClientPtr >& sleep_list )
        {
            auto p = std::make_unique< PubClientTcp >(
                context, messages, sleep_list );
            return p.release( );
        }

        inline void
        intrusive_ptr_add_ref( PubClientTcp* pc )
        {
            pc->ref_cnt_++;
            //            std::cout << "socket: " << pc->socket(
            //            ).native_handle( )
            //                      << " refcount: " << pc->ref_count( ) <<
            //                      std::endl;
        }

        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 )
                : context_{}, acceptor_{ context_,
                                         tcp::endpoint(
                                             boost::asio::ip::make_address(
                                                 address ),
                                             port ) },
                  stopping_{ false }, 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_ );
                //                std::cout << "socket " << client->socket(
                //                ).native_handle( )
                //                          << " refcount " <<
                //                          client->ref_count( )
                //                          << " accept_loop " << std::endl;
                acceptor_.async_accept(
                    client->socket( ),
                    [client,
                     this]( const boost::system::error_code& ec ) mutable {
                        //                        std::cout << "socket "
                        //                                  << client->socket(
                        //                                  ).native_handle( )
                        //                                  << " refcount " <<
                        //                                  client->ref_count( )
                        //                                  << " accept_loop
                        //                                  lambda" <<
                        //                                  std::endl;
                        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 )
                {
                    //                    std::cout << "socket " << pc->socket(
                    //                    ).native_handle( )
                    //                              << " refcount " <<
                    //                              pc->ref_count( )
                    //                              << " handle_accept" <<
                    //                              std::endl;
                    //                    std::cout << "Got a connection" <<
                    //                    std::endl;
                    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::cout << "wake up clients" << std::endl;
                for ( int i = 0; i < sleeping_clients_.size( ); ++i )
                {
                    PubClientTcp* p = nullptr;
                    {
                        if ( sleeping_clients_[ i ].get( ) == nullptr )
                        {
                            continue;
                        }
                        p = sleeping_clients_[ i ].get( );
                    }
                    //                    std::cout << i << ") socket "
                    //                              << p->socket(
                    //                              ).native_handle( ) << "
                    //                              refcount "
                    //                              << p->ref_count( ) << "\n";
                }

                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( );
                }
                // std::cout << "woke up clients" << std::endl;
                for ( int i = 0; i < sleeping_clients_.size( ); ++i )
                {
                    PubClientTcp* p = nullptr;
                    {
                        if ( sleeping_clients_[ i ].get( ) == nullptr )
                        {
                            continue;
                        }
                        p = sleeping_clients_[ i ].get( );
                    }
                    //                    std::cout << i << ") socket "
                    //                              << p->socket(
                    //                              ).native_handle( ) << "
                    //                              refcount "
                    //                              << p->ref_count( ) << "\n";
                }
            }

            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( make_message_block( key, msg, nullptr ) );
                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, msg );
                } );
            }
            boost::asio::io_context context_;
            tcp::acceptor           acceptor_;
            std::thread             thread_;
            bool                    stopping_;
            // FIXME: Check to see if this needs to be on the heap
            // with a shared ptr of some sort
            BoundedList< MessageBlock > messages_{};
            std::vector< PubClientPtr > sleeping_clients_{};
            PubDebugNotices&            debug_;
        };
    } // namespace detail
} // namespace pub_sub

#endif