/***
* ==++==
*
* Copyright (c) Microsoft Corporation. All rights reserved. 
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
* http://www.apache.org/licenses/LICENSE-2.0
* 
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
* ==--==
* =+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+
*
* send_msg_tests.cpp
*
* Tests cases for covering sending messages from websocket client.
*
* =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
****/

#include "stdafx.h"

#if defined(__cplusplus_winrt) || !defined(_M_ARM)

using namespace concurrency;
using namespace concurrency::streams;

using namespace web::websockets;
using namespace web::websockets::client;

using namespace tests::functional::websocket::utilities;

#if defined(__cplusplus_winrt)
using namespace Windows::Storage;
#endif

namespace tests { namespace functional { namespace websocket { namespace client {

SUITE(send_msg_tests)
{
utility::string_t get_full_name(const utility::string_t &name)
{
#if defined(__cplusplus_winrt)
    // On WinRT, we must compensate for the fact that we will be accessing files in the
    // Documents folder
    auto file = pplx::create_task(
        KnownFolders::DocumentsLibrary->CreateFileAsync(
            ref new Platform::String(name.c_str()), CreationCollisionOption::ReplaceExisting)).get();
    return file->Path->Data();
#else
    return name;
#endif
}

template<typename _CharType>
pplx::task<streams::streambuf<_CharType>> OPEN_R(const utility::string_t &name)
{
#if !defined(__cplusplus_winrt)
    return streams::file_buffer<_CharType>::open(name, std::ios_base::in);
#else
    auto file = pplx::create_task(
        KnownFolders::DocumentsLibrary->GetFileAsync(ref new Platform::String(name.c_str()))).get();

    return streams::file_buffer<_CharType>::open(file, std::ios_base::in);
#endif
}

// Used to prepare data for stream tests
void fill_file(const utility::string_t &name, const std::vector<uint8_t>& body, size_t repetitions = 1)
{
    std::fstream stream(get_full_name(name), std::ios_base::out | std::ios_base::trunc);

    for (size_t i = 0; i < repetitions; i++)
        stream.write((char *)&body[0], body.size());
    stream.close();
}

void fill_buffer(streams::streambuf<uint8_t> rbuf, const std::vector<uint8_t>& body, size_t repetitions = 1)
{
    size_t len = body.size();
    for (size_t i = 0; i < repetitions; i++)
        rbuf.putn_nocopy((const uint8_t *)&body[0], len).wait();
}

template<class SocketClientClass>
pplx::task<void> send_text_msg_helper(SocketClientClass& client, web::uri uri, test_websocket_server& server, const std::string& body, bool connect_client = true)
{
    server.next_message([body](test_websocket_msg msg) // Handler to verify the message sent by the client.
    {
        websocket_asserts::assert_message_equals(msg, body, test_websocket_message_type::WEB_SOCKET_UTF8_MESSAGE_TYPE);
    });

    if (connect_client)
        client.connect(uri).wait();

    websocket_outgoing_message msg;
    msg.set_utf8_message(body);
    return client.send(msg);
}

pplx::task<void> send_msg_from_stream(websocket_client& client,
                                             test_websocket_server& server,
                                             web::uri uri,
                                             const std::vector<uint8_t>& body, 
                                             streams::streambuf<uint8_t> buf,
                                             test_websocket_message_type type,
                                             bool fill_data,
                                             bool connect_client = true)
{
    server.next_message([body, type](test_websocket_msg msg)
    {
        websocket_asserts::assert_message_equals(msg, body, type);
    });

    if (connect_client)
        client.connect(uri).wait();
    if (fill_data)
        fill_buffer(buf, body);

    websocket_outgoing_message msg;
    if (type == test_websocket_message_type::WEB_SOCKET_UTF8_MESSAGE_TYPE)
        msg.set_utf8_message(streams::istream(buf), body.size());
    else if (type == test_websocket_message_type::WEB_SOCKET_BINARY_MESSAGE_TYPE)
        msg.set_binary_message(streams::istream(buf), body.size());

    return client.send(msg);
}

// Send message from input stream -> data is already populated in the stream buffer
pplx::task<void> send_msg_from_istream_helper(websocket_client& client,
                                             test_websocket_server& server,
                                             web::uri uri,
                                             const std::vector<uint8_t>& body, 
                                             streams::streambuf<uint8_t> rbuf,
                                             test_websocket_message_type type,
                                             bool connect_client = true)
{
    return send_msg_from_stream(client, server, uri, body, rbuf, type, false, connect_client);
}

pplx::task<void> send_msg_from_stream_helper(websocket_client& client,
                                             test_websocket_server& server,
                                             web::uri uri,
                                             const std::vector<uint8_t>& body, 
                                             streams::streambuf<uint8_t> rbuf,
                                             test_websocket_message_type type,
                                             bool connect_client = true)
{
    return send_msg_from_stream(client, server, uri, body, rbuf, type, true, connect_client);
}

// Send text message (no fragmentation)
TEST_FIXTURE(uri_address, send_text_msg)
{
    test_websocket_server server;
    websocket_client client;
    send_text_msg_helper(client, m_uri, server, "hello").wait();
    client.close().wait();
}

// Send text message with websocket_callback_client
TEST_FIXTURE(uri_address, send_text_msg_callback_client)
{
    test_websocket_server server;
    websocket_callback_client client;
    send_text_msg_helper(client, m_uri, server, "hello").wait();
    client.close().wait();
}

// Send text message (no fragmentation)
// Test the stream interface to send data
TEST_FIXTURE(uri_address, send_text_msg_stream)
{
    test_websocket_server server;
    streams::producer_consumer_buffer<uint8_t> rbuf;
    std::vector<uint8_t> body(26);
    memcpy(&body[0], "abcdefghijklmnopqrstuvwxyz", 26);

    websocket_client client;
    send_msg_from_stream_helper(client, server, m_uri, body, rbuf, test_websocket_message_type::WEB_SOCKET_UTF8_MESSAGE_TYPE).wait();

    rbuf.close(std::ios::out).wait();
    client.close().wait();
}

// Send Binary message (no fragmentation)
TEST_FIXTURE(uri_address, send_binary_msg)
{
    test_websocket_server server;
    streams::producer_consumer_buffer<uint8_t> rbuf;
    std::vector<uint8_t> body(6);
    memcpy(&body[0], "a\0b\0c\0", 6);

    websocket_client client;

    send_msg_from_stream_helper(client, server, m_uri, body, rbuf, test_websocket_message_type::WEB_SOCKET_BINARY_MESSAGE_TYPE).wait();
    rbuf.close(std::ios::out);
    client.close().wait();
}

// Send empty text message
// WinRT client does not handle empty messages. Verify websocket_exception is thrown.
TEST_FIXTURE(uri_address, send_empty_text_msg)
{
    test_websocket_server server;
    websocket_client client;

    client.connect(m_uri).wait();

    websocket_outgoing_message msg;
    msg.set_utf8_message("");
    VERIFY_THROWS(client.send(msg).wait(), websocket_exception);

    client.close().wait();
}

// Send multiple text messages
TEST_FIXTURE(uri_address, send_multiple_text_msges)
{
    test_websocket_server server;
    websocket_client client;

    send_text_msg_helper(client, m_uri, server, "hello1").wait();
    send_text_msg_helper(client, m_uri, server, "hello2", false).wait();

    client.close().wait();
}

// Send multiple text messages
TEST_FIXTURE(uri_address, send_multiple_text_msges_async)
{
    test_websocket_server server;
    websocket_client client;

    auto t1 = send_text_msg_helper(client, m_uri, server, "hello1");
    auto t2 = send_text_msg_helper(client, m_uri, server, "hello2", false);

    t2.wait();
    t1.wait();
    client.close().wait();
}

// Send multiple text messages from a stream
TEST_FIXTURE(uri_address, send_multiple_text_msges_stream)
{
    test_websocket_server server;
    streams::producer_consumer_buffer<uint8_t> rbuf;
    std::vector<uint8_t> body1(26);
    memcpy(&body1[0], "abcdefghijklmnopqrstuvwxyz", 26);
    std::vector<uint8_t> body2(26);
    memcpy(&body2[0], "zyxwvutsrqponmlkjihgfedcba", 26);

    websocket_client client;

    auto t1 = send_msg_from_stream_helper(client, server, m_uri, body1, rbuf, test_websocket_message_type::WEB_SOCKET_UTF8_MESSAGE_TYPE);
    auto t2 = send_msg_from_stream_helper(client, server, m_uri, body2, rbuf, test_websocket_message_type::WEB_SOCKET_UTF8_MESSAGE_TYPE, false);

    t1.wait();
    t2.wait();
    client.close().wait();
}

// Send multiple text messages from a file stream
// send uses stream::acquire API, acquire will fail for file streams.
TEST_FIXTURE(uri_address, send_text_msges_fstream)
{
    test_websocket_server server;
    utility::string_t fname = U("send_multiple_text_msges_fstream.txt");
    std::vector<uint8_t> body1(26);
    memcpy(&body1[0], "abcdefghijklmnopqrstuvwxyz", 26);
    fill_file(fname, body1, 2);
    auto file_buf = OPEN_R<uint8_t>(fname).get();
    websocket_client client;

    auto t1 = send_msg_from_istream_helper(client, server, m_uri, body1, file_buf, test_websocket_message_type::WEB_SOCKET_UTF8_MESSAGE_TYPE);
    auto t2 = send_msg_from_istream_helper(client, server, m_uri, body1, file_buf, test_websocket_message_type::WEB_SOCKET_UTF8_MESSAGE_TYPE, false);

    t1.wait();
    t2.wait();
    client.close().wait();
}

// Send multiple text messages from a container stream, where container stream has more data than what we want to send in a single message
TEST_FIXTURE(uri_address, send_text_msges_cstream)
{
    test_websocket_server server;
    std::vector<uint8_t> body(26);
    memcpy(&body[0], "abcdefghijklmnopqrstuvwxyz", 26);
    auto cbuf = streams::container_stream<std::vector<uint8_t>>::open_istream(body).streambuf();

    websocket_client client;

    auto t1 = send_msg_from_istream_helper(client, server, m_uri, std::vector<uint8_t>(body.begin(), body.begin() + body.size() / 2), cbuf, test_websocket_message_type::WEB_SOCKET_UTF8_MESSAGE_TYPE);
    auto t2 = send_msg_from_istream_helper(client, server, m_uri, std::vector<uint8_t>(body.begin() + body.size() / 2, body.end()), cbuf, test_websocket_message_type::WEB_SOCKET_UTF8_MESSAGE_TYPE, false);

    t1.wait();
    t2.wait();
    client.close().wait();
}

// Send multiple text messages from a producer consumer stream, where stream initially has less data than what we want to send in a single message
// Write data to the buffer after initiating the send, send should succeed.
TEST_FIXTURE(uri_address, send_text_msges_pcstream_lessdata)
{
    test_websocket_server server;
    streams::producer_consumer_buffer<uint8_t> rbuf;
    std::vector<uint8_t> body(26);
    memcpy(&body[0], "abcdefghijklmnopqrstuvwxyz", 26);
    fill_buffer(rbuf, body);

    server.next_message([](test_websocket_msg msg)
    {
        websocket_asserts::assert_message_equals(msg, "abcdefghijklmnopqrstuvwxyzabcd", test_websocket_message_type::WEB_SOCKET_UTF8_MESSAGE_TYPE);
    });

    websocket_client client;
    client.connect(m_uri).wait();
    websocket_outgoing_message msg;
    msg.set_utf8_message(rbuf.create_istream(), 30);
    auto t1 = client.send(msg);

    fill_buffer(rbuf, body);
    t1.wait();
    client.close().wait();
}

// Send multiple text messages from a container stream, where stream has less data than what we want to send in a single message
// Since container stream does not support in | out simultaneously, websocket send_msg will fail to read the required number of bytes 
// and throws an exception.
TEST_FIXTURE(uri_address, send_text_msges_cstream_lessdata)
{
    test_websocket_server server;
    std::vector<uint8_t> body(26);
    memcpy(&body[0], "abcdefghijklmnopqrstuvwxyz", 26);
    auto cbuf = streams::container_stream<std::vector<uint8_t>>::open_istream(body).streambuf();

    server.next_message([](test_websocket_msg msg)
    {
        websocket_asserts::assert_message_equals(msg, "abcdefghijklmnopqrstuvwxyzabcd", test_websocket_message_type::WEB_SOCKET_UTF8_MESSAGE_TYPE);
    });

    websocket_client client;
    client.connect(m_uri).wait();
    websocket_outgoing_message msg;
    msg.set_utf8_message(cbuf.create_istream(), 30);

    VERIFY_THROWS(client.send(msg).wait(), websocket_exception);
    client.close().wait();
}

// Send multiple binary messages from the same stream
TEST_FIXTURE(uri_address, send_multiple_binary_msg_same_stream)
{
    test_websocket_server server;
    streams::producer_consumer_buffer<uint8_t> rbuf;
    std::vector<uint8_t> body1(6);
    memcpy(&body1[0], "a\0b\0c\0", 6);
    std::vector<uint8_t> body2(6);
    memcpy(&body2[0], "a\0b\0c\0", 6);

    websocket_client client;

    auto t1 = send_msg_from_stream_helper(client, server, m_uri, body1, rbuf, test_websocket_message_type::WEB_SOCKET_BINARY_MESSAGE_TYPE);
    auto t2 = send_msg_from_stream_helper(client, server, m_uri, body2, rbuf, test_websocket_message_type::WEB_SOCKET_BINARY_MESSAGE_TYPE, false);

    t1.wait();
    t2.wait();
    rbuf.close(std::ios_base::out);
    client.close().wait();
}

// Send text message followed by binary message
TEST_FIXTURE(uri_address, send_text_and_binary)
{
    test_websocket_server server;
    streams::producer_consumer_buffer<uint8_t> rbuf;
    std::vector<uint8_t> body2(6);
    memcpy(&body2[0], "a\0b\0c\0", 6);

    websocket_client client;

    send_text_msg_helper(client, m_uri, server, "hello1").wait();
    send_msg_from_stream_helper(client, server, m_uri, body2, rbuf, test_websocket_message_type::WEB_SOCKET_BINARY_MESSAGE_TYPE, false).wait();

    rbuf.close(std::ios::out).wait();
    client.close().wait();
}

// Send a multi byte UTF-8 text message 
TEST_FIXTURE(uri_address, send_multi_byte_utf8_msg)
{
    test_websocket_server server;
    std::string body = "\xC3\xA0\xC3\xB8";
    websocket_client client;

    send_text_msg_helper(client, m_uri, server, body).wait();
    client.close().wait();
}

// Send a streamed text message without specifying length
TEST_FIXTURE(uri_address, send_stream_utf8_msg_no_length)
{
    test_websocket_server server;

    std::string body = "\xC3\xA0\xC3\xB8";
    std::vector<uint8_t> msgbuf(body.begin(), body.end());

    auto is = streams::container_stream<std::vector<uint8_t>>::open_istream(std::move(msgbuf));

    websocket_client client;
    {
        server.next_message([body](test_websocket_msg msg) // Handler to verify the message sent by the client.
        {
            websocket_asserts::assert_message_equals(msg, body, test_websocket_message_type::WEB_SOCKET_UTF8_MESSAGE_TYPE);
        });

        client.connect(m_uri).wait();

        websocket_outgoing_message msg;
        msg.set_utf8_message(is);
        client.send(msg).wait();
    }

    client.close().wait();
}

// Send a streamed binary message without specifying length
TEST_FIXTURE(uri_address, send_stream_binary_msg_no_length)
{
    test_websocket_server server;

    std::string body = "\x00\x01\x02\x00";
    std::vector<uint8_t> msgbuf(body.begin(), body.end());

    auto is = streams::container_stream<std::vector<uint8_t>>::open_istream(std::move(msgbuf));

    websocket_client client;
    {
        server.next_message([body](test_websocket_msg msg) // Handler to verify the message sent by the client.
        {
            websocket_asserts::assert_message_equals(msg, body, test_websocket_message_type::WEB_SOCKET_BINARY_MESSAGE_TYPE);
        });

        client.connect(m_uri).wait();

        websocket_outgoing_message msg;
        msg.set_binary_message(is);
        client.send(msg).wait();
    }

    client.close().wait();
}

} // SUITE(send_msg_tests)

}}}}

#endif