/* pub_sub subscriber implementation
Copyright (C) 2020 Jonathan Hanks

This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/
#include <cds-pubsub/sub.hh>
#include "pub_sub_intl.hh"
#include "pipeline_parser.hh"

#include <algorithm>
#include <atomic>
#include <map>

#include <boost/asio.hpp>
#include <boost/make_shared.hpp>

#include <boost/circular_buffer.hpp>

#include <cds-pubsub/sub_plugin.hh>

#include "sub_intl.hh"
#include "sub_tcp.hh"
#include "sub_udp.hh"
#include "sub_zstd.hh"
#include "make_unique.hh"

namespace pub_sub
{
    namespace detail
    {
        class TcpSubPluginApi : public pub_sub::plugins::SubscriptionPluginApi
        {
        public:
            TcpSubPluginApi( )
                : SubscriptionPluginApi( ), context_{ },
                  work_guard_{ boost::asio::make_work_guard( context_ ) },
                  thread_{ }, stopping_{ false }
            {
                thread_ = std::thread( [ this ]( ) { context_.run( ); } );
            }
            ~TcpSubPluginApi( ) override
            {
                stopping_.store( true );
                context_.stop( );
                thread_.join( );
            }

            const std::string&
            prefix( ) const override
            {
                static const std::string my_prefix = "tcp://";
                return my_prefix;
            }
            const std::string&
            version( ) const override
            {
                static const std::string my_version = "0";
                return my_version;
            }
            const std::string&
            name( ) const override
            {
                static const std::string my_name = "tcp-unicast";
                return my_name;
            }
            std::shared_ptr< pub_sub::plugins::Subscription >
            subscribe( const std::string&  conn_str,
                       SubDebugNotices&    debug_hooks,
                       pub_sub::SubHandler handler ) override
            {
                auto new_addr = detail::split_address(
                    conn_str.substr( prefix( ).size( ) ) );
                auto endpoint = detail::tcp::endpoint(
                    boost::asio::ip::make_address( new_addr.first ),
                    new_addr.second );

                return std::make_shared< SubscriptionTcp >(
                    context_,
                    endpoint,
                    std::move( handler ),
                    stopping_,
                    debug_hooks );
            }

        private:
            io_context          context_;
            work_guard          work_guard_;
            std::atomic< bool > stopping_;
            std::thread         thread_;
        };

        class UdpSubPluginApi : public pub_sub::plugins::SubscriptionPluginApi
        {
        public:
            UdpSubPluginApi( ) : SubscriptionPluginApi( )
            {
            }
            ~UdpSubPluginApi( ) override = default;

            const std::string&
            prefix( ) const override
            {
                static const std::string my_prefix = "udp://";
                return my_prefix;
            }
            const std::string&
            version( ) const override
            {
                static const std::string my_version = "0";
                return my_version;
            }
            const std::string&
            name( ) const override
            {
                static const std::string my_name = "udp-broadcast";
                return my_name;
            }
            std::shared_ptr< pub_sub::plugins::Subscription >
            subscribe( const std::string&  conn_str,
                       SubDebugNotices&    debug_hooks,
                       pub_sub::SubHandler handler ) override
            {
                auto new_addr = detail::split_address(
                    conn_str.substr( prefix( ).size( ) ) );
                auto endpoint = detail::udp::endpoint(
                    boost::asio::ip::make_address( new_addr.first ),
                    new_addr.second );

                return std::make_shared< SubscriptionUdp >(
                    endpoint, std::move( handler ), debug_hooks );
            }
        };

        class UdpMultiSubPluginApi : public UdpSubPluginApi
        {
        public:
            UdpMultiSubPluginApi( ) : UdpSubPluginApi( )
            {
            }
            ~UdpMultiSubPluginApi( ) override = default;

            const std::string&
            prefix( ) const override
            {
                const static std::string my_prefix = "multi://";
                return my_prefix;
            }

            const std::string&
            name( ) const override
            {
                const static std::string my_name = "udp-multicast";
                return my_name;
            }
            std::shared_ptr< pub_sub::plugins::Subscription >
            subscribe( const std::string&  conn_str,
                       SubDebugNotices&    debug_hooks,
                       pub_sub::SubHandler handler ) override
            {
                auto root = conn_str.substr( prefix( ).size( ) );
                auto sep = root.find( '/' );
                if ( sep == std::string::npos )
                {
                    throw std::runtime_error(
                        "Multicast connection string is "
                        "multi://local_addr/multicast addr:port" );
                }
                auto local_addr = root.substr( 0, sep );
                auto remote_addr =
                    detail::split_address( root.substr( sep + 1 ) );

                auto local_endpoint = detail::udp::endpoint(
                    boost::asio::ip::make_address( local_addr ),
                    remote_addr.second );
                auto remote_endpoint = detail::udp::endpoint(
                    boost::asio::ip::make_address( remote_addr.first ),
                    remote_addr.second );

                return std::make_shared< SubscriptionUdp >(
                    local_endpoint,
                    remote_endpoint,
                    std::move( handler ),
                    debug_hooks );
            }
        };

        class SubPluginZLibFilter
            : public pub_sub::plugins::SubscriptionPluginApi
        {
        public:
            SubPluginZLibFilter( ) : SubscriptionPluginApi( )
            {
            }
            ~SubPluginZLibFilter( ) override = default;

            const std::string&
            prefix( ) const override
            {
                static const std::string my_prefix = "zstd://";
                return my_prefix;
            }

            const std::string&
            version( ) const override
            {
                static const std::string my_version = "0";
                return my_version;
            }

            const std::string&
            name( ) const override
            {
                static const std::string my_name = "zstd compression filter";
                return my_name;
            }

            bool
            is_source( ) const override
            {
                return false;
            }

            std::shared_ptr< pub_sub::plugins::Subscription >
            subscribe( const std::string&        address,
                       pub_sub::SubDebugNotices& debug_hooks,
                       pub_sub::SubHandler       handler ) override
            {
                if ( address.find( prefix( ) ) != 0 )
                {
                    throw std::runtime_error(
                        "Invalid subscription type passed the "
                        "the zstd filter" );
                }

                return std::make_shared< detail::ZStdSubscription >(
                    debug_hooks, std::move( handler ) );
            }
        };

        /*!
         * @brief The SubscriberIntl is the logic behind the Subscriber.
         * @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 SubscriberIntl
        {
            using PluginMap = std::map<
                std::string,
                std::shared_ptr< pub_sub::plugins::SubscriptionPluginApi > >;

            using SubMap =
                std::map< SubId,
                          std::shared_ptr< pub_sub::plugins::Subscription > >;

            SubscriberIntl( )
                : debug_hooks_{ NullSubDebugNotices( ) }, plugins_{ }
            {
                auto tcp_plugin = std::make_shared< TcpSubPluginApi >( );
                plugins_.insert(
                    std::make_pair( tcp_plugin->prefix( ), tcp_plugin ) );

                auto udp_plugin = std::make_shared< UdpSubPluginApi >( );
                plugins_.insert(
                    std::make_pair( udp_plugin->prefix( ), udp_plugin ) );

                auto multi_plugin = std::make_shared< UdpMultiSubPluginApi >( );
                plugins_.insert(
                    std::make_pair( multi_plugin->prefix( ), multi_plugin ) );

                auto zstd_filter = std::make_shared< SubPluginZLibFilter >( );
                plugins_.insert(
                    std::make_pair( zstd_filter->prefix( ), zstd_filter ) );
            }

            void
            load_plugin(
                std::shared_ptr< pub_sub::plugins::SubscriptionPluginApi >
                    plugin )
            {
                std::string prefix = plugin->prefix( );

                if ( plugins_.find( prefix ) == plugins_.end( ) )
                {
                    plugins_.insert( std::make_pair( plugin->prefix( ),
                                                     std::move( plugin ) ) );
                }
            }

            void
            debug_hooks( SubDebugNotices* hooks )
            {
                debug_hooks_ =
                    ( hooks ? hooks : detail::NullSubDebugNotices( ) );
            }

            SubId
            subscribe( const std::string& address, SubHandler handler )
            {
                SubId result{ };

                auto pipeline_stages =
                    pub_sub::detail::parse_pipeline( address );
                check_pipeline_layout( pipeline_stages );

                std::reverse( pipeline_stages.begin( ),
                              pipeline_stages.end( ) );

                std::shared_ptr< pub_sub::plugins::Subscription > sub = nullptr;

                auto cur_handler = std::move( handler );
                // for each plugin
                for ( const auto& cur_stage_name : pipeline_stages )
                {
                    auto it = lookup_plugin( cur_stage_name );
                    if ( it == plugins_.end( ) )
                    {
                        throw std::runtime_error(
                            "Unable to find requested plugin" );
                    }

                    sub = it->second->subscribe( cur_stage_name,
                                                 *debug_hooks_,
                                                 std::move( cur_handler ) );

                    cur_handler = [ sub ]( SubMessage sub_message ) {
                        sub->filter( std::move( sub_message ) );
                    };
                }
                result = sub->sub_id( );
                subscriptions_.insert(
                    std::make_pair( result, std::move( sub ) ) );

                // return the last sub_id
                return result;
            }

            void
            reset_subscription( SubId subscription_id )
            {
                auto it = subscriptions_.find( subscription_id );
                if ( it != subscriptions_.end( ) )
                {
                    it->second->reset( );
                }
            }

            void
            check_pipeline_layout( const std::vector< std::string >& stages )
            {
                if ( stages.empty( ) )
                {
                    throw std::runtime_error( "Invalid connection string" );
                }
                auto require_source = true;
                for ( const auto& stage : stages )
                {
                    auto it = lookup_plugin( stage );
                    if ( it == plugins_.end( ) )
                    {
                        throw std::runtime_error(
                            "invalid connection string, unable to locate "
                            "correct plugin" );
                    }
                    if ( it->second->is_source( ) != require_source )
                    {
                        throw std::runtime_error(
                            ( require_source
                                  ? "the initial plugin in a subscription "
                                    "pipeline must be a source"
                                  : "secondary plugins in a subscription "
                                    "pipeline must be filters" ) );
                    }
                    require_source = false;
                }
            }

            PluginMap::iterator
            lookup_plugin( const std::string& stage )
            {
                return std::find_if(
                    plugins_.begin( ),
                    plugins_.end( ),
                    [ stage ]( const PluginMap::value_type& plugin ) -> bool {
                        return stage.find( plugin.second->prefix( ) ) == 0;
                    } );
            }

            SubDebugNotices* debug_hooks_;
            PluginMap        plugins_;
            SubMap           subscriptions_;
        };

    } // namespace detail

    Subscriber::Subscriber( )
        : p_{ std::make_unique< detail::SubscriberIntl >( ) }
    {
    }
    Subscriber::~Subscriber( ) = default;

    void
    Subscriber::load_plugin(
        std::shared_ptr< pub_sub::plugins::SubscriptionPluginApi > plugin )
    {
        p_->load_plugin( std::move( plugin ) );
    }

    void
    Subscriber::SetupDebugHooks( SubDebugNotices* debug_hooks )
    {
        p_->debug_hooks( debug_hooks );
    }

    SubId
    Subscriber::subscribe( const std::string& address, SubHandler handler )
    {
        return p_->subscribe( address, std::move( handler ) );
    }

    void
    Subscriber::reset_subscription( SubId subscription_id )
    {
        p_->reset_subscription( subscription_id );
    }

} // namespace pub_sub