#include <catch.hpp>

#include <pub_sub_intl.hh>
#include <pub_intl.hh>
#include <pub_tcp.hh>
#include <pub_zstd.hh>
#include <cds-pubsub/pub.hh>
#include <cds-pubsub/sub.hh>

#include <atomic>
#include <chrono>
#include <future>
#include <mutex>
#include <set>
#include <sstream>
#include <vector>

namespace
{
    class RestartSubDebug : public pub_sub::SubDebugNotices
    {
    public:
        void
        connection_terminated( const pub_sub::SubId   id,
                               const pub_sub::KeyType keyType ) override
        {
            terminate_count_++;
        }

        void
        connection_started( const pub_sub::SubId   id,
                            const pub_sub::KeyType keyType ) override
        {
            start_count_++;
        }

        std::int64_t
        start_count( ) const
        {
            return start_count_;
        }

        std::int64_t
        terminate_count( ) const
        {
            return terminate_count_;
        }

    private:
        std::int64_t start_count_{ 0 };
        std::int64_t terminate_count_{ 0 };
    };

    class SlowPubDebug : public pub_sub::PubDebugNotices
    {
    public:
        void
        slow_client_detected( const pub_sub::KeyType last_sent,
                              const pub_sub::KeyType oldest_available ) override
        {
            skip_count_++;
        }

        int
        skip_count( ) const
        {
            return skip_count_;
        }

    private:
        std::atomic< int > skip_count_{ 0 };
    };

    /*!
     * @brief Drive a depth-limited publisher against a subscriber whose
     * handler stalls on the first message, forcing the publisher-side
     * client to fall behind the retained history.
     * @returns the keys the subscriber observed, in order of arrival.
     */
    std::vector< pub_sub::KeyType >
    run_slow_subscriber( pub_sub::detail::PublisherIntlTcp& pub_tcp,
                         pub_sub::SubDebugNotices*          debug_sub,
                         const pub_sub::KeyType             last_key )
    {
        constexpr std::size_t MSG_SIZE = 4 * 1024 * 1024;

        std::mutex                      key_mutex;
        std::vector< pub_sub::KeyType > keys;

        std::promise< void >       first_msg_seen;
        std::promise< void >       gate;
        std::shared_future< void > gate_future( gate.get_future( ) );
        std::atomic< bool >        stalled_once{ false };

        std::string sub_str;
        {
            char buff[ 100 ];
            sprintf( &buff[ 0 ],
                     "tcp://127.0.0.1:%d",
                     pub_tcp.local_endpoint( ).port( ) );
            sub_str = buff;
        }

        pub_sub::Subscriber sub;
        if ( debug_sub )
        {
            sub.SetupDebugHooks( debug_sub );
        }
        sub.subscribe( sub_str, [ & ]( pub_sub::SubMessage sub_msg ) {
            if ( !stalled_once.exchange( true ) )
            {
                first_msg_seen.set_value( );
                gate_future.wait( );
            }
            std::lock_guard< std::mutex > lock( key_mutex );
            keys.push_back( sub_msg.key( ) );
        } );

        auto publish_key = [ &pub_tcp ]( pub_sub::KeyType key ) {
            pub_sub::DataPtr data( new unsigned char[ MSG_SIZE ] );
            std::fill( data.get( ),
                       data.get( ) + MSG_SIZE,
                       static_cast< unsigned char >( key ) );
            pub_tcp.publish( key, pub_sub::Message( data, MSG_SIZE ) );
        };

        /* make sure the subscriber is connected and stalled before
         * flooding it.  A subscriber that connects after a publish only
         * receives the next published message, so keep republishing key 1
         * until the subscriber reports it has seen a message. */
        auto first_msg_future = first_msg_seen.get_future( );
        auto first_msg_status = std::future_status::timeout;
        for ( int attempt = 0;
              attempt < 120 && first_msg_status != std::future_status::ready;
              ++attempt )
        {
            publish_key( 1 );
            first_msg_status =
                first_msg_future.wait_for( std::chrono::milliseconds( 250 ) );
        }
        if ( first_msg_status == std::future_status::ready )
        {
            for ( pub_sub::KeyType key = 2; key <= last_key; ++key )
            {
                publish_key( key );
            }
        }
        /* release the handler quickly so the stalled write completes well
         * inside the publisher's 2 second write deadline; always release
         * it so the subscriber can shut down on a failure */
        gate.set_value( );
        REQUIRE( first_msg_status == std::future_status::ready );

        auto deadline =
            std::chrono::steady_clock::now( ) + std::chrono::seconds( 30 );
        while ( std::chrono::steady_clock::now( ) < deadline )
        {
            {
                std::lock_guard< std::mutex > lock( key_mutex );
                if ( !keys.empty( ) && keys.back( ) == last_key )
                {
                    break;
                }
            }
            std::this_thread::sleep_for( std::chrono::milliseconds( 10 ) );
        }
        std::lock_guard< std::mutex > lock( key_mutex );
        return keys;
    }
} // namespace

TEST_CASE( "You can destroy a publisher_intl" )
{
    auto debug = pub_sub::detail::NullPubDebugNotices( );
    {
        pub_sub::detail::PublisherIntl pub( debug );
    }
    REQUIRE( true );
}

TEST_CASE( "You can destroy a publisher" )
{
    auto debug = pub_sub::detail::NullPubDebugNotices( );
    {
        pub_sub::Publisher pub( debug );
    }
    REQUIRE( true );
}

TEST_CASE( "You can send messages through tcp" )
{
    std::atomic< int > good_messages{ 0 };
    std::set< int >    messages;

    auto debug_pub = pub_sub::detail::NullPubDebugNotices( );

    pub_sub::detail::PublisherIntlTcp pub_tcp( "127.0.0.1", 0, *debug_pub );

    std::string sub_str;
    {
        std::ostringstream stream{ };
        auto               endpoint = pub_tcp.local_endpoint( );
        std::cout << "tcp://127.0.0.1:" << endpoint.port( ) << std::flush;
        stream << "tcp://127.0.0.1:" << endpoint.port( ) << std::flush;
        char buff[ 100 ];
        sprintf( &buff[ 0 ], "tcp://127.0.0.1:%d", endpoint.port( ) );
        // sub_str == stream.str();
        sub_str = buff;
        std::cout << sub_str << "\n";
        std::cout << "connecting through " << sub_str << std::endl;
    }
    auto sub = std::make_unique< pub_sub::Subscriber >( );
    sub->subscribe(
        sub_str, [ &good_messages, &messages ]( pub_sub::SubMessage sub_msg ) {
            if ( messages.find( sub_msg.key( ) ) != messages.end( ) )
            {
                return;
            }
            messages.insert( sub_msg.key( ) );
            if ( static_cast< unsigned char >( sub_msg.key( ) ) ==
                 *reinterpret_cast< unsigned char* >( sub_msg.data( ) ) )
            {
                good_messages++;
            }
        } );
    int i = 0;
    while ( good_messages.load( ) < 1000 )
    {
        pub_sub::DataPtr data( new unsigned char[ 20000 ] );
        std::fill( data.get( ),
                   data.get( ) + 20000,
                   static_cast< unsigned char >( i ) );
        pub_tcp.publish( i, pub_sub::Message( data, 20000 ) );
        std::this_thread::sleep_for( std::chrono::milliseconds( 25 ) );
        ++i;
    }
}

TEST_CASE( "You can send messages through tcp and reset the tcp connection" )
{
    std::atomic< int > good_messages{ 0 };
    std::set< int >    messages;

    RestartSubDebug debug_sub;
    auto            debug_pub = pub_sub::detail::NullPubDebugNotices( );

    pub_sub::detail::PublisherIntlTcp pub_tcp( "127.0.0.1", 0, *debug_pub );

    std::string sub_str;
    {
        std::ostringstream stream{ };
        auto               endpoint = pub_tcp.local_endpoint( );
        std::cout << "tcp://127.0.0.1:" << endpoint.port( ) << std::flush;
        stream << "tcp://127.0.0.1:" << endpoint.port( ) << std::flush;
        char buff[ 100 ];
        sprintf( &buff[ 0 ], "tcp://127.0.0.1:%d", endpoint.port( ) );
        // sub_str == stream.str();
        sub_str = buff;
        std::cout << sub_str << "\n";
        std::cout << "connecting through " << sub_str << std::endl;
    }
    auto sub = std::make_unique< pub_sub::Subscriber >( );
    sub->SetupDebugHooks( &debug_sub );
    auto sub_ptr = sub.get( );
    int  msg_count = 0;
    sub->subscribe(
        sub_str,
        [ sub_ptr, &msg_count, &good_messages, &messages ](
            pub_sub::SubMessage sub_msg ) {
            if ( msg_count % 100 == 0 )
            {
                sub_ptr->reset_subscription( sub_msg.sub_id( ) );
            }
            ++msg_count;
            if ( messages.find( sub_msg.key( ) ) != messages.end( ) )
            {
                return;
            }
            messages.insert( sub_msg.key( ) );
            if ( static_cast< unsigned char >( sub_msg.key( ) ) ==
                 *reinterpret_cast< unsigned char* >( sub_msg.data( ) ) )
            {
                good_messages++;
            }
        } );
    int i = 0;
    while ( good_messages.load( ) < 1000 )
    {
        pub_sub::DataPtr data( new unsigned char[ 20000 ] );
        std::fill( data.get( ),
                   data.get( ) + 20000,
                   static_cast< unsigned char >( i ) );
        pub_tcp.publish( i, pub_sub::Message( data, 20000 ) );
        std::this_thread::sleep_for( std::chrono::milliseconds( 25 ) );
        ++i;
    }
    REQUIRE( debug_sub.start_count( ) > 3 );
    REQUIRE( debug_sub.terminate_count( ) > 3 );
}
TEST_CASE( "A slow tcp subscriber is skipped ahead" )
{
    const pub_sub::KeyType LAST_KEY = 20;

    SlowPubDebug debug_pub;

    pub_sub::detail::PublisherIntlTcp pub_tcp(
        "127.0.0.1", 0, debug_pub, pub_sub::detail::PublisherOptions{ 3 } );

    auto keys = run_slow_subscriber( pub_tcp, nullptr, LAST_KEY );

    /* the subscriber caught back up to the live stream */
    REQUIRE( !keys.empty( ) );
    REQUIRE( keys.back( ) == LAST_KEY );
    /* but it observed a gap where it was skipped ahead */
    REQUIRE( keys.size( ) < LAST_KEY );
    /* keys always arrive in increasing order, gap or not */
    REQUIRE( std::is_sorted( keys.begin( ), keys.end( ) ) );
    REQUIRE( debug_pub.skip_count( ) >= 1 );
}
