/*
        Copyright © 2015 Gluzskiy Alexandr (sss)

    This file is part of Unknown Download Manager (UDM).

    UDM 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 2 of the License, or
    (at your option) any later version.

    UDM 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 UDM.  If not, see <http://www.gnu.org/licenses/>.

*/

#ifndef SOCKET_H_INCLUDED
#define SOCKET_H_INCLUDED

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

class socket_wraper
{
public:
    socket_wraper(boost::asio::ssl::stream<boost::asio::ip::tcp::socket>* s)
        : is_ssl(true)
    {
	socket_ssl_ = s;
    }

    socket_wraper(boost::asio::ip::tcp::socket* s)
        : is_ssl(false)
    {
	socket_ = s;
    }

    void operator=(boost::asio::ssl::stream<boost::asio::ip::tcp::socket>* s)
    {
	socket_ssl_ = s;
	is_ssl = true;
    }
    void operator=(boost::asio::ip::tcp::socket* s)
    {
	socket_ = s;
	is_ssl = false;
    }

    boost::asio::ip::tcp::socket& get_socket()
    {
	return *socket_;
    }

    boost::asio::ssl::stream<boost::asio::ip::tcp::socket>& get_ssl_socket()
    {
	return *socket_ssl_;
    }
    template <typename MutableBufferSequence> std::size_t read_some(const MutableBufferSequence& buffers)
    {
	if(is_ssl)
	    return socket_ssl_->read_some(buffers);
	else
	    return socket_->read_some(buffers);
    }
    template <typename MutableBufferSequence>
    std::size_t read_some(const MutableBufferSequence& buffers, boost::system::error_code& ec)
    {
	if(is_ssl)
	    return socket_ssl_->read_some(buffers, ec);
	else
	    return socket_->read_some(buffers, ec);
    }

    template <typename MutableBufferSequence, typename ReadHandler>
    void async_read_some(const MutableBufferSequence& buffers, ReadHandler handler)
    {
	if(is_ssl)
	    socket_ssl_->async_read_some(buffers, handler);
	else
	    socket_->async_read_some(buffers, handler);
    }

    template <typename ConstBufferSequence, typename WriteHandler>
    void async_write_some(const ConstBufferSequence& buffers, WriteHandler handler)
    {
	if(is_ssl)
	    socket_ssl_->async_write_some(buffers, handler);
	else
	    socket_->async_write_some(buffers, handler);
    }
    ~socket_wraper()
    {
	if(is_ssl)
	    delete socket_ssl_;
	else
	    delete socket_;
    }

private:
    bool is_ssl;
    boost::asio::ip::tcp::socket* socket_ = nullptr;
    boost::asio::ssl::stream<boost::asio::ip::tcp::socket>* socket_ssl_ = nullptr;
};

#endif