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

#ifndef PUB_SUB2_PUB_UDP_HH
#define PUB_SUB2_PUB_UDP_HH

#include "pub_intl.hh"
#include <time.h>

namespace pub_sub
{
    namespace detail
    {
        struct PublisherIntlUdp : public plugins::PublisherInstance
        {
            PublisherIntlUdp( const std::string& local_address,
                              const std::string& remote_address,
                              int                port,
                              PubDebugNotices&   debug )
                : context_{ }, local_endpoint_{ boost::asio::ip::make_address(
                                                    local_address ),
                                                0 },
                  remote_endpoint_{ boost::asio::ip::make_address(
                                        remote_address ),
                                    static_cast< unsigned short >( port ) },
                  s_{ context_ }, retry_s_{ context_, local_endpoint_ },
                  work_{ boost::asio::make_work_guard( context_ ) }, thread_{ },
                  recent_retries_( 10 ), resend_req_buffer_{ }, debug_{ debug }
            {
                std::cout << "local: " << local_endpoint_ << std::endl;
                std::cout << "remote: " << remote_endpoint_ << std::endl;
                s_.open( local_endpoint_.protocol( ) );
                s_.set_option( udp::socket::broadcast( true ) );
                s_.bind( local_endpoint_ );
                std::cout << "Bound to " << s_.local_endpoint( ) << std::endl;
                s_.connect( remote_endpoint_ );

                std::cout << "Retry bound to "
                          << retry_s_.local_endpoint( ).address( ) << ":"
                          << retry_s_.local_endpoint( ).port( ) << std::endl;

                start_read_loop( );

                thread_ = std::thread( [ this ]( ) { context_.run( ); } );
            }
            PublisherIntlUdp( const PublisherIntlUdp& ) = delete;
            PublisherIntlUdp( PublisherIntlUdp&& ) = delete;
            PublisherIntlUdp& operator=( const PublisherIntlUdp& ) = delete;
            PublisherIntlUdp& operator=( PublisherIntlUdp&& ) = delete;
            ~PublisherIntlUdp( ) override = default;

            void
            publish( KeyType key, Message msg ) override
            {
                auto msg_block = make_message_block( key, msg, nullptr );
                context_.post(
                    [ this, msg_block ]( ) { send_message( msg_block ); } );
            }

            void
            send_partial_message( MessageBlockPtr        msg,
                                  UdpHeader::block_index start_block,
                                  UdpHeader::size_type   block_count )
            {
                if ( !msg || block_count < 1 )
                {
                    return;
                }

                auto blocks = static_cast< UdpHeader::block_index >(
                    msg->msg( ).length / UdpHeader::MAX_CHUNK( ) );
                auto remainder = msg->msg( ).length % UdpHeader::MAX_CHUNK( );
                if ( remainder )
                {
                    ++blocks;
                }

                if ( start_block >= blocks )
                {
                    return;
                }

                std::size_t initial_offset =
                    start_block * UdpHeader::MAX_CHUNK( );
                char* cur_data =
                    reinterpret_cast< char* >( msg->msg( ).data( ) ) +
                    initial_offset;
                std::size_t remaining_data =
                    msg->msg( ).length - initial_offset;

                UdpHeader::size_type final_block_size =
                    remaining_data % UdpHeader::MAX_CHUNK( );
                UdpHeader::block_index block_index{ start_block };

                auto sender_address = static_cast< std::uint32_t >(
                    local_endpoint_.address( ).to_v4( ).to_uint( ) );
                auto sender_port =
                    static_cast< std::uint16_t >( local_endpoint_.port( ) );
                while ( remaining_data && block_count )
                {
                    --block_count;

                    UdpHeader header{ };
                    header.key = msg->key( );
                    header.block_number = block_index;
                    header.block_count = blocks;
                    header.final_block_size = final_block_size;
                    header.sender_address = sender_address;
                    header.sender_port = sender_port;
                    header.flags = static_cast< std::uint8_t >(
                        UdpBlockType::RETRANSMISSION );

                    std::size_t block_size =
                        ( header.is_last_block( ) ? final_block_size
                                                  : UdpHeader::MAX_CHUNK( ) );
                    std::array< boost::asio::const_buffer, 2 > buffers;
                    buffers[ 0 ] =
                        boost::asio::const_buffer( &header, sizeof( header ) );
                    buffers[ 1 ] =
                        boost::asio::const_buffer( cur_data, block_size );

                    s_.send( buffers );
                    remaining_data -= block_size;
                    cur_data += block_size;
                    ++block_index;
                }

                debug_.retransmit_sent( msg->key( ), block_count );
            }

            void
            send_message( MessageBlockPtr msg )
            {
                auto blocks = static_cast< UdpHeader::block_index >(
                    msg->msg( ).length / UdpHeader::MAX_CHUNK( ) );
                auto remainder = msg->msg( ).length % UdpHeader::MAX_CHUNK( );
                if ( remainder )
                {
                    ++blocks;
                }
                char* cur_data =
                    reinterpret_cast< char* >( msg->msg( ).data( ) );
                std::size_t remaining_data = msg->msg( ).length;

                UdpHeader::size_type final_block_size =
                    remaining_data % UdpHeader::MAX_CHUNK( );
                UdpHeader::block_index block_index{ 0 };

                auto sender_address = static_cast< std::uint32_t >(
                    // local_endpoint_.address( ).to_v4( ).to_uint( ) );
                    retry_s_.local_endpoint( ).address( ).to_v4( ).to_uint( ) );
                auto sender_port = static_cast< std::uint16_t >(
                    // local_endpoint_.port( )
                    retry_s_.local_endpoint( ).port( ) );
                timespec sleeptime{ 0, 100000 };
                while ( remaining_data )
                {
                    UdpHeader header{ };
                    header.key = msg->key( );
                    header.block_number = block_index;
                    header.block_count = blocks;
                    header.final_block_size = final_block_size;
                    header.sender_address = sender_address;
                    header.sender_port = sender_port;
                    header.flags =
                        static_cast< std::uint8_t >( UdpBlockType::DATA );

                    std::size_t block_size =
                        ( header.is_last_block( ) ? final_block_size
                                                  : UdpHeader::MAX_CHUNK( ) );
                    std::array< boost::asio::const_buffer, 2 > buffers;
                    buffers[ 0 ] =
                        boost::asio::const_buffer( &header, sizeof( header ) );
                    buffers[ 1 ] =
                        boost::asio::const_buffer( cur_data, block_size );

                    if ( !( drop_counter_ % 10 == 0 && drop_counter_ ) )
                    {
                        nanosleep( &sleeptime, nullptr );
                        s_.send( buffers );
                    }
                    drop_counter_ = 0;
                    timespec ts{ }, ts_rem{ };
                    ts.tv_sec = 0;
                    ts.tv_nsec = 2000;
                    nanosleep( &ts, &ts_rem );
                    // drop_counter_++;
                    remaining_data -= block_size;
                    cur_data += block_size;
                    ++block_index;
                }
                messages_.push( std::move( msg ) );
            }

            void
            start_read_loop( )
            {
                context_.post( [ this ]( ) { read_loop( ); } );
            }

            void
            read_loop( )
            {
                // std::cout << "Entering read loop" << std::endl;
                boost::asio::mutable_buffer buffer(
                    &resend_req_buffer_, sizeof( resend_req_buffer_ ) );
                retry_s_.async_receive_from(
                    buffer,
                    resend_req_endpoint_,
                    [ this ]( const boost::system::error_code& ec,
                              std::size_t bytes_transferred ) {
                        if ( ec ||
                             bytes_transferred != sizeof( UdpResendRequest ) )
                        {
                            // std::cout << "read loop error ec = " << ec
                            //          << " bytes = " << bytes_transferred
                            //          << std::endl;
                            start_read_loop( );
                            return;
                        }
                        debug_.retransmit_request_seen(
                            resend_req_buffer_.msg_key,
                            resend_req_buffer_.block_count );
                        //                        std::cout << "Retransmit
                        //                        request seen on "
                        //                                  <<
                        //                                  resend_req_buffer_.msg_key
                        //                                  << " "
                        //                                  <<
                        //                                  resend_req_buffer_.start_block
                        //                                  << " "
                        //                                  <<
                        //                                  resend_req_buffer_.block_count
                        //                                  << std::endl;
                        request_retry( );
                    } );
            }

            void
            request_retry( )
            {
                register_retry( );
                start_read_loop( );
            }
            void
            register_retry( )
            {
                if ( retry_request_seen( resend_req_buffer_ ) )
                {
                    return;
                }
                MessageBlockPtr msg =
                    find_old_message( resend_req_buffer_.msg_key );
                if ( !msg )
                {
                    return;
                }
                // std::cout << "Retransmit request seen on " << msg->key() << "
                // " << resend_req_buffer_.start_block << " " <<
                // resend_req_buffer_.block_count << std::endl;

                context_.post( [ this,
                                 message = std::move( msg ),
                                 start = resend_req_buffer_.start_block,
                                 count = resend_req_buffer_.block_count ]( ) {
                    send_partial_message( message, start, count );
                } );
            }

            bool
            retry_request_seen( UdpResendRequest& retry )
            {
                auto it = std::find_if(
                    recent_retries_.begin( ),
                    recent_retries_.end( ),
                    [ &retry ]( const UdpResendRequest& cur ) -> bool {
                        return ( retry == cur );
                    } );
                if ( it != recent_retries_.end( ) )
                {
                    return true;
                }
                recent_retries_.push_front( retry );
                return false;
            }

            MessageBlockPtr
            find_old_message( KeyType key )
            {
                auto cur = messages_.tail( ).get( );
                while ( cur )
                {
                    if ( cur->key( ) == key )
                    {
                        return cur;
                    }
                    cur = cur->next( ).get( );
                }
                return nullptr;
            }

            boost::asio::io_context                    context_;
            udp::endpoint                              local_endpoint_;
            udp::endpoint                              remote_endpoint_;
            udp::socket                                s_;
            udp::socket                                retry_s_;
            BoundedList< MessageBlock >                messages_{ };
            work_guard                                 work_;
            std::thread                                thread_;
            boost::circular_buffer< UdpResendRequest > recent_retries_;

            boost::asio::ip::udp::endpoint resend_req_endpoint_;
            PubDebugNotices&               debug_;
            int                            drop_counter_{ 0 };
            UdpResendRequest               resend_req_buffer_;
        };
    } // namespace detail
} // namespace pub_sub
#endif // PUB_SUB2_PUB_UDP_HH
