/*
 * Copyright (c) 2011 Tobias Markmann
 * Licensed under the simplified BSD license.
 * See Documentation/Licenses/BSD-simplified.txt for more information.
 */

/*
 * Copyright (c) 2015-2016 Isode Limited.
 * All rights reserved.
 * See the COPYING file for more information.
 */

#include <boost/bind.hpp>
#include <boost/random/mersenne_twister.hpp>
#include <boost/random/uniform_int.hpp>
#include <boost/random/variate_generator.hpp>
#include <boost/smart_ptr/make_shared.hpp>

#include <QA/Checker/IO.h>

#include <cppunit/extensions/HelperMacros.h>
#include <cppunit/extensions/TestFactoryRegistry.h>

#include <Swiften/Base/Algorithm.h>
#include <Swiften/Base/ByteArray.h>
#include <Swiften/Base/Concat.h>
#include <Swiften/Base/Log.h>
#include <Swiften/Base/StartStopper.h>
#include <Swiften/Crypto/CryptoProvider.h>
#include <Swiften/Crypto/PlatformCryptoProvider.h>
#include <Swiften/EventLoop/DummyEventLoop.h>
#include <Swiften/FileTransfer/ByteArrayReadBytestream.h>
#include <Swiften/FileTransfer/ByteArrayWriteBytestream.h>
#include <Swiften/FileTransfer/SOCKS5BytestreamClientSession.h>
#include <Swiften/FileTransfer/SOCKS5BytestreamRegistry.h>
#include <Swiften/JID/JID.h>
#include <Swiften/Network/DummyConnection.h>
#include <Swiften/Network/DummyTimerFactory.h>
#include <Swiften/StringCodecs/Hexify.h>

using namespace Swift;

static boost::mt19937 randomGen;

class SOCKS5BytestreamClientSessionTest : public CppUnit::TestFixture {
    CPPUNIT_TEST_SUITE(SOCKS5BytestreamClientSessionTest);
    CPPUNIT_TEST(testForSessionReady);
    CPPUNIT_TEST(testErrorHandlingHello);
    CPPUNIT_TEST(testErrorHandlingRequest);
    CPPUNIT_TEST(testWriteBytestream);
    CPPUNIT_TEST(testReadBytestream);
    CPPUNIT_TEST_SUITE_END();

public:
    SOCKS5BytestreamClientSessionTest() : destinationAddressPort(HostAddressPort(HostAddress("127.0.0.1"), 8888)) {}

    void setUp() {
        crypto = boost::shared_ptr<CryptoProvider>(PlatformCryptoProvider::create());
        destination = "092a44d859d19c9eed676b551ee80025903351c2";
        randomGen.seed(static_cast<unsigned int>(time(nullptr)));
        eventLoop = new DummyEventLoop();
        timerFactory = new DummyTimerFactory();
        connection = boost::make_shared<MockeryConnection>(failingPorts, true, eventLoop);
        //connection->onDataSent.connect(boost::bind(&SOCKS5BytestreamServerSessionTest::handleDataWritten, this, _1));
        //stream1 = boost::make_shared<ByteArrayReadBytestream>(createByteArray("abcdefg")));
//        connection->onDataRead.connect(boost::bind(&SOCKS5BytestreamClientSessionTest::handleDataRead, this, _1));
    }

    void tearDown() {
        //connection.reset();
        delete timerFactory;
        delete eventLoop;
    }

    void testForSessionReady() {
        TestHelper helper;
        connection->onDataSent.connect(boost::bind(&TestHelper::handleConnectionDataWritten, &helper, _1));

        SOCKS5BytestreamClientSession::ref clientSession = boost::make_shared<SOCKS5BytestreamClientSession>(connection, destinationAddressPort, destination, timerFactory);
        clientSession->onSessionReady.connect(boost::bind(&TestHelper::handleSessionReady, &helper, _1));

        clientSession->start();
        eventLoop->processEvents();
        CPPUNIT_ASSERT(createByteArray("\x05\x01\x00", 3) == helper.unprocessedInput);

        helper.unprocessedInput.clear();
        serverRespondHelloOK();
        eventLoop->processEvents();
        CPPUNIT_ASSERT_EQUAL(createByteArray("\x05\x01\x00\x03", 4), createByteArray(&helper.unprocessedInput[0], 4));
        CPPUNIT_ASSERT_EQUAL(createByteArray(static_cast<char>(destination.size())), createByteArray(static_cast<char>(helper.unprocessedInput[4])));
        CPPUNIT_ASSERT_EQUAL(createByteArray(destination), createByteArray(&helper.unprocessedInput[5], destination.size()));
        CPPUNIT_ASSERT_EQUAL(createByteArray("\x00", 1), createByteArray(&helper.unprocessedInput[5 + destination.size()], 1));

        helper.unprocessedInput.clear();
        serverRespondRequestOK();
        eventLoop->processEvents();
        CPPUNIT_ASSERT_EQUAL(true, helper.sessionReadyCalled);
        CPPUNIT_ASSERT_EQUAL(false, helper.sessionReadyError);
    }

    void testErrorHandlingHello() {
        TestHelper helper;
        connection->onDataSent.connect(boost::bind(&TestHelper::handleConnectionDataWritten, &helper, _1));

        SOCKS5BytestreamClientSession::ref clientSession = boost::make_shared<SOCKS5BytestreamClientSession>(connection, destinationAddressPort, destination, timerFactory);
        clientSession->onSessionReady.connect(boost::bind(&TestHelper::handleSessionReady, &helper, _1));

        clientSession->start();
        eventLoop->processEvents();
        CPPUNIT_ASSERT_EQUAL(createByteArray("\x05\x01\x00", 3), helper.unprocessedInput);

        helper.unprocessedInput.clear();
        serverRespondHelloAuthFail();
        eventLoop->processEvents();

        CPPUNIT_ASSERT_EQUAL(true, helper.sessionReadyCalled);
        CPPUNIT_ASSERT_EQUAL(true, helper.sessionReadyError);
        CPPUNIT_ASSERT_EQUAL(true, connection->disconnectCalled);
    }

    void testErrorHandlingRequest() {
        TestHelper helper;
        connection->onDataSent.connect(boost::bind(&TestHelper::handleConnectionDataWritten, &helper, _1));

        SOCKS5BytestreamClientSession::ref clientSession = boost::make_shared<SOCKS5BytestreamClientSession>(connection, destinationAddressPort, destination, timerFactory);
        clientSession->onSessionReady.connect(boost::bind(&TestHelper::handleSessionReady, &helper, _1));

        clientSession->start();
        eventLoop->processEvents();
        CPPUNIT_ASSERT_EQUAL(createByteArray("\x05\x01\x00", 3), helper.unprocessedInput);

        helper.unprocessedInput.clear();
        serverRespondHelloOK();
        eventLoop->processEvents();
        CPPUNIT_ASSERT_EQUAL(createByteArray("\x05\x01\x00\x03", 4), createByteArray(&helper.unprocessedInput[0], 4));
        CPPUNIT_ASSERT_EQUAL(createByteArray(static_cast<char>(destination.size())), createByteArray(static_cast<char>(helper.unprocessedInput[4])));
        CPPUNIT_ASSERT_EQUAL(createByteArray(destination), createByteArray(&helper.unprocessedInput[5], destination.size()));
        CPPUNIT_ASSERT_EQUAL(createByteArray("\x00", 1), createByteArray(&helper.unprocessedInput[5 + destination.size()], 1));

        helper.unprocessedInput.clear();
        serverRespondRequestFail();
        eventLoop->processEvents();
        CPPUNIT_ASSERT_EQUAL(true, helper.sessionReadyCalled);
        CPPUNIT_ASSERT_EQUAL(true, helper.sessionReadyError);
        CPPUNIT_ASSERT_EQUAL(true, connection->disconnectCalled);
    }

    void testWriteBytestream() {
        TestHelper helper;
        connection->onDataSent.connect(boost::bind(&TestHelper::handleConnectionDataWritten, &helper, _1));

        SOCKS5BytestreamClientSession::ref clientSession = boost::make_shared<SOCKS5BytestreamClientSession>(connection, destinationAddressPort, destination, timerFactory);
        clientSession->onSessionReady.connect(boost::bind(&TestHelper::handleSessionReady, &helper, _1));

        clientSession->start();
        eventLoop->processEvents();

        helper.unprocessedInput.clear();
        serverRespondHelloOK();
        eventLoop->processEvents();

        helper.unprocessedInput.clear();
        serverRespondRequestOK();
        eventLoop->processEvents();
        CPPUNIT_ASSERT_EQUAL(true, helper.sessionReadyCalled);
        CPPUNIT_ASSERT_EQUAL(false, helper.sessionReadyError);

        boost::shared_ptr<ByteArrayWriteBytestream> output = boost::make_shared<ByteArrayWriteBytestream>();
        clientSession->startReceiving(output);

        ByteArray transferData = generateRandomByteArray(1024);
        connection->onDataRead(createSafeByteArrayRef(vecptr(transferData), transferData.size()));
        CPPUNIT_ASSERT_EQUAL(transferData, output->getData());
    }

    void testReadBytestream() {
        TestHelper helper;
        connection->onDataSent.connect(boost::bind(&TestHelper::handleConnectionDataWritten, &helper, _1));

        SOCKS5BytestreamClientSession::ref clientSession = boost::make_shared<SOCKS5BytestreamClientSession>(connection, destinationAddressPort, destination, timerFactory);
        clientSession->onSessionReady.connect(boost::bind(&TestHelper::handleSessionReady, &helper, _1));

        clientSession->start();
        eventLoop->processEvents();

        helper.unprocessedInput.clear();
        serverRespondHelloOK();
        eventLoop->processEvents();

        helper.unprocessedInput.clear();
        serverRespondRequestOK();
        eventLoop->processEvents();
        CPPUNIT_ASSERT_EQUAL(true, helper.sessionReadyCalled);
        CPPUNIT_ASSERT_EQUAL(false, helper.sessionReadyError);

        helper.unprocessedInput.clear();
        ByteArray transferData = generateRandomByteArray(1024);
        boost::shared_ptr<ByteArrayReadBytestream> input = boost::make_shared<ByteArrayReadBytestream>(transferData);
        clientSession->startSending(input);
        eventLoop->processEvents();

        CPPUNIT_ASSERT_EQUAL(createByteArray(vecptr(transferData), transferData.size()), helper.unprocessedInput);
    }



private:
    static ByteArray generateRandomByteArray(size_t len) {
        boost::uniform_int<> dist(0, 255);
        boost::variate_generator<boost::mt19937&, boost::uniform_int<> > randomByte(randomGen, dist);
        ByteArray result(len);
        for (size_t i=0; i < len; ++i ) {
            result[i] = static_cast<unsigned char>(randomByte());
        }
        return result;
    }

    // Server responses
    void serverRespondHelloOK() {
        connection->onDataRead(createSafeByteArrayRef("\x05\00", 2));
    }

    void serverRespondHelloAuthFail() {
        connection->onDataRead(createSafeByteArrayRef("\x05\xFF", 2));
    }

    void serverRespondRequestOK() {
        boost::shared_ptr<SafeByteArray> dataToSend = createSafeByteArrayRef("\x05\x00\x00\x03", 4);
        append(*dataToSend, createSafeByteArray(static_cast<char>(destination.size())));
        append(*dataToSend, createSafeByteArray(destination));
        append(*dataToSend, createSafeByteArray("\x00", 1));
        connection->onDataRead(dataToSend);
    }

    void serverRespondRequestFail() {
        boost::shared_ptr<SafeByteArray> correctData = createSafeByteArrayRef("\x05\x00\x00\x03", 4);
        append(*correctData, createSafeByteArray(static_cast<char>(destination.size())));
        append(*correctData, createSafeByteArray(destination));
        append(*correctData, createSafeByteArray("\x00", 1));

        boost::shared_ptr<SafeByteArray> dataToSend;
        //ByteArray failingData = Hexify::unhexify("8417947d1d305c72c11520ea7d2c6e787396705e72c312c6ccc3f66613d7cae1b91b7ab48e8b59a17d559c15fb51");
        //append(dataToSend, failingData);
        //SWIFT_LOG(debug) << "hexed: " << Hexify::hexify(failingData) << std::endl;
        do {
            ByteArray rndArray = generateRandomByteArray(correctData->size());
            dataToSend = createSafeByteArrayRef(vecptr(rndArray), rndArray.size());
        } while (*dataToSend == *correctData);
        connection->onDataRead(dataToSend);
    }

private:
    struct TestHelper {
        TestHelper() : sessionReadyCalled(false), sessionReadyError(false) {}
        ByteArray unprocessedInput;
        bool sessionReadyCalled;
        bool sessionReadyError;

        void handleConnectionDataWritten(const SafeByteArray& data) {
            append(unprocessedInput, data);
            //SWIFT_LOG(debug) << "unprocessedInput (" << unprocessedInput.size() <<  "): " << Hexify::hexify(unprocessedInput) << std::endl;
        }

        void handleSessionReady(bool error) {
            sessionReadyCalled = true;
            sessionReadyError = error;
        }
    };


private:
    struct MockeryConnection : public Connection, public EventOwner, public boost::enable_shared_from_this<MockeryConnection> {
        public:
        MockeryConnection(const std::vector<HostAddressPort>& failingPorts, bool isResponsive, EventLoop* eventLoop) : eventLoop(eventLoop), failingPorts(failingPorts), isResponsive(isResponsive), disconnectCalled(false) {}

            void listen() { assert(false); }
            void connect(const HostAddressPort& address) {
                hostAddressPort = address;
                if (isResponsive) {
                    bool fail = std::find(failingPorts.begin(), failingPorts.end(), address) != failingPorts.end();
                    eventLoop->postEvent(boost::bind(boost::ref(onConnectFinished), fail));
                }
            }

            HostAddressPort getLocalAddress() const { return HostAddressPort(); }

            HostAddressPort getRemoteAddress() const { return HostAddressPort(); }

            void disconnect() {
                disconnectCalled = true;
            }

            void write(const SafeByteArray& data) {
                eventLoop->postEvent(boost::ref(onDataWritten), shared_from_this());
                onDataSent(data);
            }

            boost::signal<void (const SafeByteArray&)> onDataSent;

            EventLoop* eventLoop;
            boost::optional<HostAddressPort> hostAddressPort;
            std::vector<HostAddressPort> failingPorts;
            bool isResponsive;
            bool disconnectCalled;
    };

private:
    HostAddressPort destinationAddressPort;
    std::string destination;
    DummyEventLoop* eventLoop;
    DummyTimerFactory* timerFactory;
    boost::shared_ptr<MockeryConnection> connection;
    const std::vector<HostAddressPort> failingPorts;
    boost::shared_ptr<CryptoProvider> crypto;
};

CPPUNIT_TEST_SUITE_REGISTRATION(SOCKS5BytestreamClientSessionTest);