/***
* ==++==
*
* 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.
*
* ==--==
* =+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+
*
* Defines a test client to handle requests and sending responses.
*
* =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
****/

#include "stdafx.h"

#include "cpprest/uri.h"
#include "cpprest/details/http_helpers.h"
#include "test_http_client.h"
#ifdef _WIN32
#include <winhttp.h>
#pragma warning ( push )
#pragma warning ( disable : 4457 )
#include <agents.h>
#pragma warning ( pop )
#endif

using namespace web; using namespace utility;

namespace tests { namespace functional { namespace http { namespace utilities {

// Flatten the http_headers into a name:value pairs separated by a carriage return and line feed.
utility::string_t flatten_http_headers(const std::map<utility::string_t, utility::string_t> &headers)
{
    utility::string_t flattened_headers;
    for(auto iter = headers.begin(); iter != headers.end(); ++iter)
    {
        utility::string_t temp((*iter).first + U(":") + (*iter).second + U("\r\n"));
        flattened_headers.append(utility::string_t(temp.begin(), temp.end()));
    }
    return flattened_headers;
}

#ifdef _WIN32

// Helper function to query for the size of header values.
static void query_header_length(HINTERNET request_handle, DWORD header, DWORD &length)
{
    WinHttpQueryHeaders(
        request_handle,
        header,
        WINHTTP_HEADER_NAME_BY_INDEX,
        WINHTTP_NO_OUTPUT_BUFFER,
        &length,
        WINHTTP_NO_HEADER_INDEX);
}

// Helper function to get the status code from a WinHTTP response.
static void parse_status_code(HINTERNET request_handle, unsigned short &code)
{
    DWORD length = 0;
    query_header_length(request_handle, WINHTTP_QUERY_STATUS_CODE, length);
    utility::string_t buffer;
    buffer.resize(length);
    WinHttpQueryHeaders(
        request_handle,
        WINHTTP_QUERY_STATUS_CODE,
        WINHTTP_HEADER_NAME_BY_INDEX,
        &buffer[0],
        &length,
        WINHTTP_NO_HEADER_INDEX);
    code = (unsigned short)_wtoi(buffer.c_str());
}

// Helper function to trim leading and trailing null characters from a string.
static void trim_nulls(utility::string_t &str)
{
    size_t index;
    for(index = 0; index < str.size() && str[index] == 0; ++index);
    str.erase(0, index);
    index;
    for(index = str.size(); index > 0 && str[index - 1] == 0; --index);
    str.erase(index);
}

// Helper function to get the reason phrase from a WinHTTP response.
static void parse_reason_phrase(HINTERNET request_handle, utility::string_t &phrase)
{
    DWORD length = 0;
    query_header_length(request_handle, WINHTTP_QUERY_STATUS_TEXT, length);
    phrase.resize(length);
    WinHttpQueryHeaders(
        request_handle,
        WINHTTP_QUERY_STATUS_TEXT,
        WINHTTP_HEADER_NAME_BY_INDEX,
        &phrase[0],
        &length,
        WINHTTP_NO_HEADER_INDEX);
    // WinHTTP reports back the wrong length, trim any null characters.
    trim_nulls(phrase);
}

/// <summary>
/// Parses a string containing Http headers.
/// </summary>
static void parse_winhttp_headers(HINTERNET request_handle, utf16char *headersStr, test_response *p_response)
{
    // Status code and reason phrase.
    parse_status_code(request_handle, p_response->m_status_code);
    parse_reason_phrase(request_handle, p_response->m_reason_phrase);

    utf16char *context = nullptr;
    utf16char *line = wcstok_s(headersStr, U("\r\n"), &context);
    while(line != nullptr)
    {
        const utility::string_t header_line(line);
        const size_t colonIndex = header_line.find_first_of(U(":"));
        if(colonIndex != utility::string_t::npos)
        {
            utility::string_t key = header_line.substr(0, colonIndex);
            utility::string_t value = header_line.substr(colonIndex + 1, header_line.length() - colonIndex - 1);
            web::http::details::trim_whitespace(key);
            web::http::details::trim_whitespace(value);
            p_response->m_headers[key] = value;
        }
        line = wcstok_s(nullptr, U("\r\n"), &context);
    }
}

class _test_http_client
{
public:
    _test_http_client(const utility::string_t &uri) 
        : m_uri(uri), m_hSession(nullptr), m_hConnection(nullptr)
    {}

    unsigned long open()
    {
        // Open session.
        m_hSession = WinHttpOpen(
            U("test_http_client"),
            WINHTTP_ACCESS_TYPE_DEFAULT_PROXY,
            WINHTTP_NO_PROXY_NAME,
            WINHTTP_NO_PROXY_BYPASS,
            WINHTTP_FLAG_ASYNC);
        if(!m_hSession)
        {
            return GetLastError();
        }

        
        // Set timeouts.
        int multiplier=10;
        if(!WinHttpSetTimeouts(m_hSession, 60000*multiplier, 60000*multiplier, 30000*multiplier, 30000*multiplier))
        {
            return GetLastError();
        }

        // Set max connection to use per server to 1.
        DWORD maxConnections = 1;
        if(!WinHttpSetOption(m_hSession, WINHTTP_OPTION_MAX_CONNS_PER_SERVER, &maxConnections, sizeof(maxConnections)))
        {
            return GetLastError();
        }

        // Register asynchronous callback.
        if(WINHTTP_INVALID_STATUS_CALLBACK == WinHttpSetStatusCallback(
            m_hSession,
            &_test_http_client::completion_callback,
            WINHTTP_CALLBACK_FLAG_ALL_COMPLETIONS,
            0))
        {
            return GetLastError();
        }

        // Open connection.
        ::http::uri u(m_uri);
        unsigned int port = u.is_port_default() ? INTERNET_DEFAULT_PORT : u.port();
        m_hConnection = WinHttpConnect(
            m_hSession,
            u.host().c_str(),
            (INTERNET_PORT)port,
            0);
        if(m_hConnection == nullptr)
        {
            return GetLastError();
        }
        return 0;
    }

    unsigned long close()
    {
        // Release memory for each request.
        std::for_each(m_responses_memory.begin(), m_responses_memory.end(), [](test_response *p_response)
        {
            delete p_response;
        });

        if(m_hConnection != nullptr)
        {
            if(WinHttpCloseHandle(m_hConnection) == NULL)
            {
                return GetLastError();
            }
        }

        if(m_hSession != nullptr)
        {
            // Unregister the callback.
            if(!WinHttpSetStatusCallback(
                m_hSession,
                NULL,
                WINHTTP_CALLBACK_FLAG_ALL_NOTIFICATIONS,
                NULL))
            {
                return GetLastError();
            }

            if(WinHttpCloseHandle(m_hSession) == NULL)
            {
                return GetLastError();
            }
        }
        return 0;
    }

    unsigned long request(
        const utility::string_t &method,
        const utility::string_t &path,
        const std::map<utility::string_t, utility::string_t> headers,
        void * data,
        size_t data_length)
    {
        HINTERNET request_handle = WinHttpOpenRequest(
            m_hConnection,
            method.c_str(),
            path.c_str(),
            NULL,
            WINHTTP_NO_REFERER,
            WINHTTP_DEFAULT_ACCEPT_TYPES,
            0);
        if(request_handle == nullptr)
        {
            return GetLastError();
        }

        // Add headers.
        if(!headers.empty())
        {
            utility::string_t flattened_headers = flatten_http_headers(headers);
            if(!WinHttpAddRequestHeaders(
                request_handle,
                flattened_headers.c_str(),
                (DWORD)flattened_headers.length(),
                WINHTTP_ADDREQ_FLAG_ADD))
            {
                return GetLastError();
            }
        }

        if(!WinHttpSendRequest(
            request_handle,
            WINHTTP_NO_ADDITIONAL_HEADERS,
            0,
            data,
            (DWORD)data_length,
            (DWORD)data_length,
            (DWORD_PTR)new test_response(this)))
        {
            return GetLastError();
        }
        return 0;
    }

    test_response * wait_for_response()
    {
        return wait_for_responses(1)[0];
    }
    
    pplx::task<test_response *> next_response()
    {
        return pplx::create_task([this]() -> test_response *
        {
            return wait_for_response();
        });
    }
    
    std::vector<test_response *> wait_for_responses(const size_t count)
    {
        std::vector<test_response *> m_test_responses;
        for(size_t i = 0; i < count; ++i)
        {
            m_test_responses.push_back(Concurrency::receive(m_responses));
        }
        return m_test_responses;
    }
    
    std::vector<pplx::task<test_response *>> next_responses(const size_t count)
    {
        std::vector<pplx::task_completion_event<test_response *>> events;
        std::vector<pplx::task<test_response *>> responses;
        for(size_t i = 0; i < count; ++i)
        {
            events.push_back(pplx::task_completion_event<test_response *>());
            responses.push_back(pplx::create_task(events[i]));
        }
        pplx::create_task([this, count, events]()
        {
            for(size_t i = 0; i < count; ++i)
            {
                events[i].set(wait_for_response());
            }
        });
        return responses;
    }

private:

    // WinHTTP callback.
    static void CALLBACK completion_callback(
        HINTERNET hRequestHandle,
        DWORD_PTR context,
        DWORD statusCode,
        void* statusInfo,
        DWORD)
    {
        test_response * p_response = reinterpret_cast<test_response *>(context);
        if(p_response != nullptr)
        {
            if(statusCode == WINHTTP_CALLBACK_STATUS_REQUEST_ERROR)
            {
                WINHTTP_ASYNC_RESULT *pStatusInfo = static_cast<WINHTTP_ASYNC_RESULT *>(statusInfo);
                pStatusInfo;
                throw std::exception("Error in WinHTTP callback");
            }
            else if(statusCode == WINHTTP_CALLBACK_STATUS_SENDREQUEST_COMPLETE)
            {
                if(!WinHttpReceiveResponse(hRequestHandle, NULL))
                {
                    throw std::exception("Error receiving response");
                }
            }
            else if(statusCode == WINHTTP_CALLBACK_STATUS_HEADERS_AVAILABLE)
            {
                DWORD headers_length;
                WinHttpQueryHeaders(
                    hRequestHandle,
                    WINHTTP_QUERY_RAW_HEADERS_CRLF,
                    WINHTTP_HEADER_NAME_BY_INDEX,
                    WINHTTP_NO_OUTPUT_BUFFER,
                    &headers_length,
                    WINHTTP_NO_HEADER_INDEX);

                // Now allocate buffer for headers and query for them.
                std::vector<unsigned char> header_raw_buffer;
                header_raw_buffer.resize(headers_length);
                utf16char * header_buffer = reinterpret_cast<utf16char *>(&header_raw_buffer[0]);
                if(!WinHttpQueryHeaders(
                    hRequestHandle,
                    WINHTTP_QUERY_RAW_HEADERS_CRLF,
                    WINHTTP_HEADER_NAME_BY_INDEX,
                    header_buffer,
                    &headers_length,
                    WINHTTP_NO_HEADER_INDEX))
                {
                    throw std::exception("Error querying for headers");
                }
                parse_winhttp_headers(hRequestHandle, header_buffer, p_response);

                // Check to see if the response has a body or not.
                if(!WinHttpQueryDataAvailable(hRequestHandle, nullptr))
                {
                    throw std::exception("Error reading response body");
                }
            }
            else if(statusCode == WINHTTP_CALLBACK_STATUS_DATA_AVAILABLE)
            {
                DWORD num_bytes = *(PDWORD)statusInfo;

                if ( num_bytes > 0 )
                {
                    size_t current_size =  p_response->m_data.size();
                    p_response->m_data.resize(current_size+(size_t)num_bytes);
                    
                    // Actual WinHTTP call to read in body.
                    if(!WinHttpReadData(
                        hRequestHandle,
                        &p_response->m_data[current_size],
                        (DWORD)num_bytes,
                        NULL))
                    {
                        throw std::exception("Error reading response body");
                    }
                }
                else
                {
                    WinHttpCloseHandle(hRequestHandle);
                    p_response->m_client->m_responses_memory.push_back(p_response);
                    Concurrency::asend(p_response->m_client->m_responses, p_response);
                }
            }
            else if(statusCode == WINHTTP_CALLBACK_STATUS_READ_COMPLETE)
            {
                if(!WinHttpQueryDataAvailable(hRequestHandle, nullptr))
                {
                    throw std::exception("Error reading response body");
                }
            }
        }
    }

    Concurrency::unbounded_buffer<test_response *> m_responses;

    // Used to store all requests to simplify memory management.
    std::vector<test_response *> m_responses_memory;

    const utility::string_t m_uri;
    HINTERNET m_hSession;
    HINTERNET m_hConnection;
};
#else
class _test_http_client
{
private:
    const web::http::uri m_uri;
    typename web::http::client::http_client m_client;
    std::vector<pplx::task<web::http::http_response>> m_responses;
    std::vector<test_response*> m_test_responses;
    
public:
    _test_http_client(utility::string_t uri) : m_uri(web::http::uri::encode_uri(uri)), m_client(m_uri.authority()) {}

    unsigned long open()  { return 0; }
    unsigned long close() { return 0; }

    unsigned long request(
        const utility::string_t &method,
        const utility::string_t &path,
        const std::map<utility::string_t, utility::string_t>& headers,
        void * data,
        size_t data_length)
    {
        auto localHeaders = headers;
        localHeaders["User-Agent"] = "test_http_client";
        web::http::http_request request;
        request.set_method(method);
        request.set_request_uri(web::http::uri_builder(m_uri).append_path(path).to_uri());
        auto& hDest = request.headers();
        for (auto it = localHeaders.begin(); it != localHeaders.end(); ++it)
        {
            auto& currentValue = hDest[it->first];
            if (currentValue.empty())
                currentValue = it->second;
            else
                currentValue = currentValue + U(", ") + it->second;
        }
        request.set_body(utility::string_t(reinterpret_cast<const char*>(data), data_length));

        m_responses.push_back(m_client.request(request));
        return 0;
    }

    test_response * wait_for_response()
    {
        return wait_for_responses(1)[0];
    }
    
    pplx::task<test_response *> next_response()
    {
        return pplx::create_task([this]() -> test_response *
        {
            return wait_for_response();
        });
    }
    
    std::vector<test_response *> wait_for_responses(const size_t count)
    {
        if (count > m_responses.size())
            throw std::logic_error("count too big");

        std::vector<test_response *> m_test_responses;
        for(size_t i = 0; i < count; ++i)
        {
            auto response = m_responses[0].get();

            auto tr = new test_response(this);
            tr->m_status_code = response.status_code();
            tr->m_reason_phrase = response.reason_phrase();
            for (auto it = response.headers().begin(); it != response.headers().end(); ++it)
            {
                tr->m_headers[it->first] = it->second;
            }
            tr->m_data = response.extract_vector().get();

            m_test_responses.push_back(tr);
            m_responses.erase(m_responses.begin());
        }
        return m_test_responses;
    }
    
    std::vector<pplx::task<test_response *>> next_responses(const size_t count)
    {
        std::vector<pplx::task<test_response *>> result;
        for (size_t i = 0; i < count; ++i)
        {
            result.push_back(next_response());
        }
        return result;
    }
};
#endif

test_http_client::test_http_client(const web::http::uri &uri)
{
    m_impl = std::unique_ptr<_test_http_client>(new _test_http_client(uri.to_string()));
}

test_http_client::~test_http_client()
{
}

test_http_client::test_http_client(test_http_client &&other) : m_impl(std::move(other.m_impl)) {}

test_http_client & test_http_client::operator=(test_http_client &&other)
{
    if(this != &other)
    {
        this->m_impl = std::move(other.m_impl);
    }
    return *this;
}

unsigned long test_http_client::open() { return m_impl->open(); }
unsigned long test_http_client::close() { return m_impl->close(); }

unsigned long test_http_client::request(const utility::string_t &method, const utility::string_t &path)
{
    return request(method, path, std::map<utility::string_t, utility::string_t>());
}
unsigned long test_http_client::request(
        const utility::string_t &method, 
        const utility::string_t &path, 
        const std::map<utility::string_t, utility::string_t> &headers)
{
    return request(method, path, headers, std::string());
}
unsigned long test_http_client::request(
        const utility::string_t &method, 
        const utility::string_t &path, 
        const std::string &data)
{
    return request(method, path, std::map<utility::string_t, utility::string_t>(), data);
}
unsigned long test_http_client::request(
        const utility::string_t &method, 
        const utility::string_t &path,
        const utility::string_t &content_type,
        const std::string &data)
{
    std::map<utility::string_t, utility::string_t> headers;
    headers[U("Content-Type")] = content_type;
    return request(method, path, headers, data);
}

unsigned long test_http_client::request(
        const utility::string_t &method, 
        const utility::string_t &path, 
        const std::map<utility::string_t, utility::string_t> &headers,
        const std::string &data)
{
    return m_impl->request(method, path, headers, (void *)&data[0], data.size());
}

test_response * test_http_client::wait_for_response() { return m_impl->wait_for_response(); }
pplx::task<test_response *> test_http_client::next_response() { return m_impl->next_response(); }
std::vector<test_response *> test_http_client::wait_for_responses(const size_t count) { return m_impl->wait_for_responses(count); }
std::vector<pplx::task<test_response *>> test_http_client::next_responses(const size_t count) { return m_impl->next_responses(count); }

}}}}