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

#ifndef PUB_SUB2_SUB_UDP_HH
#define PUB_SUB2_SUB_UDP_HH

#include <cds-pubsub/sub_plugin.hh>
#include "sub_intl.hh"
#include <poll.h>

namespace pub_sub
{
    namespace detail
    {
        enum AddBlockType
        {
            VALID_BLOCK,
            INVALID_BLOCK,
            DUPLICATE_BLOCK,
        };

        struct UdpMessageBlock
        {
            UdpMessageBlock( const UdpHeader& header,
                             unsigned char*   data,
                             const TimePoint& now )
                : header_{ header }, block_count_{ 0 }, end_block_{ 0 },
                  mask_( header_.block_count, 0 ),
                  data_{ std::make_unique< unsigned char[] >(
                      header_.message_size( ) ) },
                  needs_retrans_{ false },
                  forward_progress_made_( max_progress( ) ),
                  start_time_{ now }, end_time_{ now }
            {
                add_block( header, data, now );
            }

            bool
            complete( ) const noexcept
            {
                return block_count_ == header_.block_count;
            }

            AddBlockType
            add_block( const UdpHeader& header,
                       unsigned char*   data,
                       const TimePoint& now )
            {
                if ( mask_[ header.block_number ] )
                {
                    return DUPLICATE_BLOCK;
                }
                if ( header.block_number >= header_.block_count ||
                     header.key != header_.key )
                {
                    return INVALID_BLOCK;
                }

                std::copy( data,
                           data + header.payload_size( ),
                           data_.get( ) + header.offset_in_message( ) );
                ++block_count_;
                end_time_ = now;
                mask_[ header.block_number ] = 1;
                if ( header.block_number > end_block_ )
                {
                    needs_retrans_ = true;
                }
                if ( header.block_number >= end_block_ )
                {
                    forward_progress_made_ = max_progress( );
                }
                end_block_ = std::max( static_cast< UdpHeader::block_index >(
                                           header.block_number + 1 ),
                                       end_block_ );

                // We only need to review the history of mask on a retransmit
                if ( header.flags ==
                     static_cast< std::uint8_t >(
                         UdpBlockType::RETRANSMISSION ) )
                {
                    auto end = mask_.begin( ) + end_block_;
                    auto it = std::find( mask_.begin( ), end, 0 );
                    needs_retrans_ = ( it != end );
                }
                return VALID_BLOCK;
            }

            void
            age_progress_marker( )
            {
                if ( !complete( ) && forward_progress_made_ > 0 )
                {
                    --forward_progress_made_;
                }
            }

            bool
            forward_progress_due( ) const
            {
                return !complete( ) && forward_progress_made_ <= 0;
            }

            bool
            need_retransmit( ) const noexcept
            {
                return needs_retrans_ || !forward_progress_made_;
            }

            KeyType
            key( ) const
            {
                return header_.key;
            }

            std::size_t
            size( ) const
            {
                return header_.message_size( );
            }

            DataPtr
            extract_data( )
            {
                header_.block_count = 0;
                header_.block_number = 0;
                header_.final_block_size = 0;
                header_.key = NEXT_PUB_MSG( );
                return std::move( data_ );
            }

            UdpHeader::block_index
            next_block_at_head( ) const noexcept
            {
                return end_block_;
            }

            static constexpr int
            max_progress( )
            {
                return 5;
            }

            MilliSec
            receive_duration( ) const
            {
                return std::chrono::duration_cast< MilliSec >( end_time_ -
                                                               start_time_ );
            }

            UdpHeader                    header_;
            UdpHeader::block_index       block_count_;
            UdpHeader::block_index       end_block_;
            std::vector< unsigned char > mask_;
            DataPtr                      data_;
            TimePoint                    start_time_;
            TimePoint                    end_time_;

            bool needs_retrans_;
            int  forward_progress_made_;
        };

        class UdpBuildCache
        {
        public:
            UdpBuildCache( ) : in_process_{ nullptr, nullptr, nullptr, nullptr }
            {
            }
            void
            clear_empty( )
            {
                auto it = std::find_if(
                    in_process_.begin( ),
                    in_process_.end( ),
                    []( const std::unique_ptr< UdpMessageBlock >& p ) -> bool {
                        return ( p && p->key( ) == NEXT_PUB_MSG( ) );
                    } );
                if ( it != in_process_.end( ) )
                {
                    *it = nullptr;
                }
            }

            UdpMessageBlock*
            find_message_add_block( const UdpHeader& header,
                                    unsigned char*   data,
                                    SubDebugNotices& debug,
                                    SubId            debug_sub_id,
                                    const TimePoint& now )
            {
                auto msg =
                    find_message( header, data, debug, debug_sub_id, now );
                if ( msg->add_block( header, data, now ) == DUPLICATE_BLOCK )
                {
                    debug.duplicate_packet( debug_sub_id );
                }
                return msg;
            }

            UdpMessageBlock*
            find_message( const UdpHeader& header,
                          unsigned char*   data,
                          SubDebugNotices& debug,
                          SubId            debug_sub_id,
                          TimePoint        now )
            {
                auto it = std::find_if(
                    in_process_.begin( ),
                    in_process_.end( ),
                    [&header](
                        const std::unique_ptr< UdpMessageBlock >& p ) -> bool {
                        if ( p )
                        {
                            return p->key( ) == header.key;
                        }
                        return false;
                    } );
                if ( it != in_process_.end( ) )
                {
                    return it->get( );
                }

                KeyType lowest_key = NEXT_PUB_MSG( );
                auto    lowest_it = in_process_.begin( );
                for ( it = in_process_.begin( ); it != in_process_.end( );
                      ++it )
                {
                    if ( *it == nullptr )
                    {
                        lowest_it = it;
                        break;
                    }
                    if ( ( *it )->key( ) < lowest_key )
                    {
                        lowest_key = ( *it )->key( );
                        lowest_it = it;
                    }
                }
                debug.message_started( debug_sub_id, header.key );
                if ( *lowest_it )
                {
                    debug.message_dropped( debug_sub_id,
                                           ( *lowest_it )->key( ) );
                }
                *lowest_it =
                    std::make_unique< UdpMessageBlock >( header, data, now );
                return lowest_it->get( );
            }

            std::string
            retransmit_string( ) const
            {
                std::string s;
                s.reserve( in_process_.size( ) );
                for ( const auto& entry : in_process_ )
                {
                    if ( !entry || !entry->need_retransmit( ) )
                    {
                        s.append( " " );
                    }
                    else
                    {
                        s.append( "+" );
                    }
                }
                return s;
            }

            template < typename Callback >
            void
            retransmit_request( Callback&& cb )
            {
                for ( const auto& entry : in_process_ )
                {
                    if ( !entry )
                    {
                        continue;
                    }
                    if ( !entry->need_retransmit( ) )
                    {
                        continue;
                    }
                    UdpHeader::block_index end_index = entry->end_block_;
                    auto&                  mask = entry->mask_;

                    end_index = std::min(
                        end_index,
                        static_cast< UdpHeader::block_index >( mask.size( ) ) );

                    UdpHeader::block_index i{ 0 };
                    while ( i < end_index )
                    {
                        // skip stretches of data that we have
                        while ( i < end_index && mask[ i ] )
                        {
                            ++i;
                        }
                        if ( i < end_index )
                        {
                            // find the end of this stretch of missing data
                            UdpHeader::block_index j = i;
                            while ( j < end_index && !mask[ j ] )
                            {
                                ++j;
                            }
                            cb( *entry, i, j - i );
                            // past the end of the missing block
                            i = j;
                        }
                    }
                    if ( entry->forward_progress_due( ) )
                    {
                        cb( *entry, entry->next_block_at_head( ), 5 );
                    }
                }
            }

            void
            age_progress_markers( )
            {
                for ( auto& entry : in_process_ )
                {
                    if ( !entry )
                    {
                        continue;
                    }
                    entry->age_progress_marker( );
                }
            }

        private:
            std::array< std::unique_ptr< UdpMessageBlock >, 4 > in_process_;
        };

        class SubscriptionUdp : public pub_sub::plugins::Subscription
        {
        public:
            /*!
             * @brief Regular udp listener
             * @param endpoint The local interface to listen on
             * @param handler
             */
            SubscriptionUdp( udp::endpoint    endpoint,
                             SubHandler       handler,
                             SubDebugNotices& debug )
                : Subscription( ), context_{}, s_{ context_, endpoint },
                  endpoint_{ endpoint }, handler_{ std::move( handler ) },
                  debug_{ debug },
                  work_{ boost::asio::make_work_guard( context_ ) }, thread_{}
            {
                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( ) < 10 * 1024 * 1024 )
                {
                    option = 10 * 1024 * 1024;
                    s_.set_option( option );
                    s_.get_option( option );
                    std::cout
                        << "Socket recieve buffer size is: " << option.value( )
                        << std::endl;
                }
                // s_.open(endpoint_.protocol());
                // s_.bind(endpoint_);
                context_.post( [this]( ) { read_loop( ); } );
                thread_ = std::thread( [this]( ) { context_.run( ); } );
            }
            /*!
             * @brief The multicast udp subscription constructor
             * @param local_endpoint Local endpoint to listen on
             * @param remote_endpoint Multi cast address to listen on
             * @param handler
             */
            SubscriptionUdp( udp::endpoint    local_endpoint,
                             udp::endpoint    remote_endpoint,
                             SubHandler       handler,
                             SubDebugNotices& debug )
                : Subscription( ), context_{}, s_{ context_ },
                  endpoint_{ remote_endpoint },
                  handler_{ std::move( handler ) }, debug_{ debug },
                  work_{ boost::asio::make_work_guard( context_ ) }, thread_{}
            {
                /* reference points that helped get this working on linux.
                 * http://boost.2283326.n4.nabble.com/No-packets-received-with-boost-asio-ip-udp-socket-set-to-multicast-group-td2585161.html
                 * https://stackoverflow.com/questions/8675623/boostasioipmulticastjoin-group-does-not-work
                 */
                s_.open( local_endpoint.protocol( ) );
                s_.bind( remote_endpoint );
                s_.set_option( boost::asio::ip::multicast::join_group(
                    remote_endpoint.address( ).to_v4( ),
                    local_endpoint.address( ).to_v4( ) ) );

                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( ) < 10 * 1024 * 1024 )
                {
                    option = 10 * 1024 * 1024;
                    s_.set_option( option );
                    s_.get_option( option );
                    std::cout
                        << "Socket recieve buffer size is: " << option.value( )
                        << std::endl;
                }
                // s_.open(endpoint_.protocol());
                // s_.bind(endpoint_);
                context_.post( [this]( ) { read_loop( ); } );
                thread_ = std::thread( [this]( ) { context_.run( ); } );
            }
            ~SubscriptionUdp( ) override = default;

        private:
            using packet_payload_t =
                std::array< unsigned char, UdpHeader::MAX_CHUNK( ) >;

            using used_key_lookup_t = boost::circular_buffer< KeyType >;

            void
            find_start_point( UdpHeader&        local_header,
                              packet_payload_t& local_data )
            {
                std::array< boost::asio::mutable_buffer, 2 > buffers{};
                buffers[ 0 ] = boost::asio::mutable_buffer(
                    &local_header, sizeof( local_header ) );
                buffers[ 1 ] = boost::asio::mutable_buffer(
                    local_data.data( ), local_data.size( ) );
                udp::endpoint endpoint;
                while ( true )
                {
                    auto size = s_.receive_from( buffers, endpoint );
                    if ( size <= sizeof( local_header ) ||
                         !local_header.valid( ) )
                    {
                        continue;
                    }
                    if ( local_header.flags ==
                         static_cast< std::uint8_t >(
                             UdpBlockType::RETRANSMISSION ) )
                    {
                        continue;
                    }
                    auto latest_acceptable =
                        std::min( local_header.block_count / 2, 10 );
                    if ( local_header.block_number <= latest_acceptable )
                    {
                        break;
                    }
                }
            }

            void
            handle_message( UdpMessageBlock&   msg,
                            used_key_lookup_t& processed_keys,
                            UdpBuildCache&     build_cache )
            {
                KeyType     key = msg.key( );
                std::size_t size = msg.size( );
                auto        duration = msg.receive_duration( );

                SubMessage sub_msg(
                    sub_id( ), key, Message( msg.extract_data( ), size ) );
                build_cache.clear_empty( );
                processed_keys.push_back( key );
                handler_( std::move( sub_msg ) );
                debug_.message_received( sub_id( ), key, size, duration );
            }

            bool
            from_processed_message( used_key_lookup_t& processed_keys,
                                    KeyType            key ) const noexcept
            {
                // duplicate of a finished message
                return std::find( processed_keys.begin( ),
                                  processed_keys.end( ),
                                  key ) != processed_keys.end( );
            }

            template < typename BuffersType, typename Rep, typename Period >
            std::size_t
            read_with_timeout(
                BuffersType                                 buffers,
                const std::chrono::duration< Rep, Period >& timeout,
                boost::system::error_code&                  ec )
            {
                std::size_t size = 0;
                ::pollfd    p_info{};
                p_info.fd = s_.native_handle( );
                p_info.events = POLLIN;
                p_info.revents = 0;

                auto timeout_ =
                    std::chrono::duration_cast< std::chrono::milliseconds >(
                        timeout );

                bool is_timeout{ false };
                do
                {
                    ec.clear( );
                    auto rc = poll( &p_info, 1, timeout_.count( ) );
                    if ( rc == 0 )
                    {
                        is_timeout = true;
                        break;
                    }
                    else
                    {
                        if ( rc < 1 )
                        {
                            auto err = errno;
                            if ( err == EINTR || err == EAGAIN )
                            {
                                continue;
                            }
                        }
                        udp::endpoint remote_endpoint{};
                        size =
                            s_.receive_from( buffers, remote_endpoint, 0, ec );
                        break;
                    }
                } while ( true );
                if ( is_timeout )
                {
                    ec = boost::asio::error::timed_out;
                    size = 0;
                }
                return size;
            }

            void
            request_retransmits( const udp::endpoint& source,
                                 UdpBuildCache&       build_cache )
            {
                build_cache.retransmit_request(
                    [this, &source]( UdpMessageBlock&       msg,
                                     UdpHeader::block_index start,
                                     UdpHeader::size_type   count ) {
                        UdpResendRequest req{};
                        req.msg_key = msg.key( );
                        req.sender_id =
                            s_.local_endpoint( ).address( ).to_v4( ).to_uint( );
                        req.req_number = 1;
                        req.start_block = start;
                        req.block_count = count;

                        s_.send_to( boost::asio::buffer( &req, sizeof( req ) ),
                                    source );
                        debug_.retransmit_request_made(
                            sub_id( ), req.msg_key, count );
                    } );
            }

            [[noreturn]] void
            read_loop( )
            {
                UdpHeader        local_header{};
                packet_payload_t local_data{};

                used_key_lookup_t processed_keys( 25 );
                UdpBuildCache     in_process{};

                find_start_point( local_header, local_data );

                auto now = std::chrono::steady_clock::now( );

                auto cur_msg = in_process.find_message_add_block(
                    local_header, local_data.data( ), debug_, sub_id( ), now );
                if ( cur_msg->complete( ) )
                {
                    handle_message( *cur_msg, processed_keys, in_process );
                }

                auto source_addr = boost::asio::ip::make_address_v4(
                    local_header.sender_address );
                udp::endpoint source_endpoint =
                    udp::endpoint( source_addr, local_header.sender_port );

                const std::chrono::milliseconds retransmit_wait{ 10 };
                const std::chrono::milliseconds read_timeout{ 2 };
                auto check_retransmit_at = now + retransmit_wait;

                while ( true )
                {
                    std::array< boost::asio::mutable_buffer, 2 > buffers{};
                    buffers[ 0 ] = boost::asio::mutable_buffer(
                        &local_header, sizeof( local_header ) );
                    buffers[ 1 ] = boost::asio::mutable_buffer(
                        local_data.data( ), local_data.size( ) );
                    udp::endpoint remote_endpoint;

                    boost::system::error_code ec{};
                    auto size = read_with_timeout( buffers, read_timeout, ec );
                    if ( ec != boost::asio::error::timed_out )
                    {

                        if ( size <= sizeof( local_header ) ||
                             !local_header.valid( ) )
                        {
                            continue;
                        }
                        //                        if ( local_header.flags ==
                        //                             static_cast< std::uint8_t
                        //                             >(
                        //                                 UdpBlockType::RETRANSMISSION
                        //                                 ) )
                        //                        {
                        //                            std::cout << "Retran block
                        //                            found "
                        //                                      <<
                        //                                      local_header.key
                        //                                      << " "
                        //                                      <<
                        //                                      local_header.block_number
                        //                                      << std::endl;
                        //                        }

                        if ( from_processed_message( processed_keys,
                                                     local_header.key ) )
                        {
                            continue;
                        }

                        cur_msg = in_process.find_message_add_block(
                            local_header,
                            local_data.data( ),
                            debug_,
                            sub_id( ),
                            now );

                        int sample = (int)( local_data[ 0 ] );
                        //                        std::cout << "Received " <<
                        //                        size << " bytes "
                        //                                  << local_header.key
                        //                                  << " #"
                        //                                  <<
                        //                                  local_header.block_number
                        //                                  << "/"
                        //                                  <<
                        //                                  local_header.block_count
                        //                                  << " ("
                        //                                  <<
                        //                                  local_header.payload_size(
                        //                                  ) << ") "
                        //                                  << sample << " "
                        //                                  <<
                        //                                  in_process.retransmit_string(
                        //                                  )
                        //                                  << std::endl;

                        if ( cur_msg->complete( ) )
                        {
                            handle_message(
                                *cur_msg, processed_keys, in_process );
                        }
                    }
                    in_process.age_progress_markers( );
                    now = std::chrono::steady_clock::now( );
                    if ( now > check_retransmit_at )
                    {
                        std::string pre = in_process.retransmit_string( );
                        request_retransmits( source_endpoint, in_process );
                        //                        std::cout << "Retransmit check
                        //                        ... " << pre << ":"
                        //                                  <<
                        //                                  in_process.retransmit_string(
                        //                                  )
                        //                                  << std::endl;
                        check_retransmit_at = now + retransmit_wait;
                    }
                }
            }

            io_context       context_;
            udp::endpoint    endpoint_;
            udp::socket      s_;
            SubHandler       handler_;
            SubDebugNotices& debug_;
            work_guard       work_;
            std::thread      thread_;
        };

    } // namespace detail
} // namespace pub_sub

#endif // PUB_SUB2_SUB_UDP_HH
