#include "Win32WebServer.h"

#include "FileLogger.h"
#include "StringUtils.h"

Win32WebServer::Win32WebServer()
	: mListenUrl("http://localhost:80/49152")
{

}

Win32WebServer::~Win32WebServer()
{
	clear();
}

void Win32WebServer::clear()
{
	// Remove url
	::HttpRemoveUrl(mWorkingQueue, StringUtils::convert(mListenUrl).c_str());

	// Close the Request Queue handle.
	if (mWorkingQueue)
	{
		::CloseHandle(mWorkingQueue);
	}

	// Call HttpTerminate.
	::HttpTerminate(HTTP_INITIALIZE_SERVER, nullptr);
}

void Win32WebServer::initialize()
{
	HTTPAPI_VERSION HttpApiVersion = HTTPAPI_VERSION_1;

    auto retCode = ::HttpInitialize(HttpApiVersion, HTTP_INITIALIZE_SERVER, nullptr);
	if (FAILED(retCode))
	{
		MLOG_ERROR("Failed to initialize win32 http server.");
		return;
	}

	retCode = ::HttpCreateHttpHandle(&mWorkingQueue, 0);
	if (FAILED(retCode))
	{
		MLOG_ERROR("Failed to create request queue handle.");
		return;
	}

	retCode = ::HttpAddUrl(mWorkingQueue, StringUtils::convert(mListenUrl).c_str(), nullptr);
	if (FAILED(retCode))
	{
		MLOG_ERROR("Failed to register queue to url: " << mListenUrl);
		return;
	}
}

void Win32WebServer::onRequest(PHTTP_REQUEST request)
{
    const std::wstring wide_url = request->CookedUrl.pFullUrl;
    const auto url = StringUtils::convert(wide_url);

    switch (request->Verb)
    {
    case HttpVerbGET:
        MLOG_INFO("Got a GET request for: " << url);
        auto result = sendHttpResponse(request, 200, "OK", "Hey! You hit the server \r\n");
        break;
    case HttpVerbPOST:
        MLOG_INFO("Got a POST request for: " << url);
        result = sendHttpPostResponse(request);
        break;

    default:
        MLOG_INFO("Got an unknown request for: " << url);
        result = sendHttpResponse(request, 503, "Not Implemented");
        break;
    }
}

void Win32WebServer::initializeHttpResponse(HTTP_RESPONSE& response, USHORT StatusCode, const std::string& reason)
{
    RtlZeroMemory((&response), sizeof(*(&response)));
    response.StatusCode = StatusCode;
    response.pReason = reason.c_str();
    response.ReasonLength = (USHORT)strlen(reason.c_str());
}

#define MAX_ULONG_STR ((ULONG) sizeof("4294967295"))

DWORD Win32WebServer::sendHttpPostResponse(PHTTP_REQUEST pRequest)
{
    DWORD           result;
    auto hTempFile = INVALID_HANDLE_VALUE;

    // Allocate space for an entity buffer. Buffer can be increased on demand.
    ULONG EntityBufferLength = 2048;
    PUCHAR pEntityBuffer = reinterpret_cast<PUCHAR>(::HeapAlloc(::GetProcessHeap(), 0, (EntityBufferLength)));
    if (pEntityBuffer == nullptr)
    {
        result = ERROR_NOT_ENOUGH_MEMORY;
        MLOG_ERROR("Insufficient resources");
        return result;
    }

    HTTP_RESPONSE response;
    initializeHttpResponse(response, 200, "OK");

    if (pRequest->Flags & HTTP_REQUEST_FLAG_MORE_ENTITY_BODY_EXISTS)
    {
        // The entity body is sent over multiple calls. Collect these in a file and send back. Create a temporary file.
        TCHAR  szTempName[MAX_PATH + 1];
        if (::GetTempFileName(L".", L"New", 0, szTempName) == 0)
        {
            result = ::GetLastError();
            MLOG_ERROR("Failed to set up temp file for buffer with: " << result);
            if (pEntityBuffer)
            {
                HeapFree(GetProcessHeap(), 0, (pEntityBuffer));
            }
            return result;
        }

        hTempFile = CreateFile(szTempName, GENERIC_READ | GENERIC_WRITE,
            0,                  // Do not share.
            NULL,               // No security descriptor.
            CREATE_ALWAYS,      // Overrwrite existing.
            FILE_ATTRIBUTE_NORMAL,    // Normal file.
            NULL
        );
        if (hTempFile == INVALID_HANDLE_VALUE)
        {
            result = GetLastError();
            MLOG_ERROR("Failed to create temp file for buffer with: " << result);
            if (pEntityBuffer)
            {
                HeapFree(GetProcessHeap(), 0, (pEntityBuffer));
            }
            return result;
        }

        DWORD           bytesSent;
        ULONG           TempFileBytesWritten;
        CHAR            szContentLength[MAX_ULONG_STR];
        HTTP_DATA_CHUNK dataChunk;

        ULONG TotalBytesRead = 0;
        ULONG BytesRead = 0;
        do
        {
            // Read the entity chunk from the request.
            BytesRead = 0;
            result = ::HttpReceiveRequestEntityBody(mWorkingQueue, pRequest->RequestId, 0,
                pEntityBuffer, EntityBufferLength, &BytesRead, NULL);

            switch (result)
            {
            case NO_ERROR:
                if (BytesRead != 0)
                {
                    TotalBytesRead += BytesRead;
                    ::WriteFile(hTempFile, pEntityBuffer, BytesRead, &TempFileBytesWritten, NULL);
                }
                break;
            case ERROR_HANDLE_EOF:

                //
                // The last request entity body has been read.
                // Send back a response. 
                //
                // To illustrate entity sends via 
                // HttpSendResponseEntityBody, the response will 
                // be sent over multiple calls. To do this,
                // pass the HTTP_SEND_RESPONSE_FLAG_MORE_DATA
                // flag.

                if (BytesRead != 0)
                {
                    TotalBytesRead += BytesRead;
                    WriteFile(
                        hTempFile,
                        pEntityBuffer,
                        BytesRead,
                        &TempFileBytesWritten,
                        NULL
                    );
                }

                //
                // Because the response is sent over multiple
                // API calls, add a content-length.
                //
                // Alternatively, the response could have been
                // sent using chunked transfer encoding, by  
                // passimg "Transfer-Encoding: Chunked".
                //

                // NOTE: Because the TotalBytesread in a ULONG
                //       are accumulated, this will not work
                //       for entity bodies larger than 4 GB. 
                //       For support of large entity bodies,
                //       use a ULONGLONG.
                // 


                sprintf_s(szContentLength, MAX_ULONG_STR, "%lu", TotalBytesRead);

                ADD_KNOWN_HEADER(
                    response,
                    HttpHeaderContentLength,
                    szContentLength
                );

                result =
                    HttpSendHttpResponse(
                        hReqQueue,           // ReqQueueHandle
                        pRequest->RequestId, // Request ID
                        HTTP_SEND_RESPONSE_FLAG_MORE_DATA,
                        &response,       // HTTP response
                        NULL,            // pReserved1
                        &bytesSent,      // bytes sent-optional
                        NULL,            // pReserved2
                        0,               // Reserved3
                        NULL,            // LPOVERLAPPED
                        NULL             // pReserved4
                    );

                if (result != NO_ERROR)
                {
                    wprintf(
                        L"HttpSendHttpResponse failed with %lu \n",
                        result
                    );
                    goto Done;
                }

                //
                // Send entity body from a file handle.
                //
                dataChunk.DataChunkType =
                    HttpDataChunkFromFileHandle;

                dataChunk.FromFileHandle.
                    ByteRange.StartingOffset.QuadPart = 0;

                dataChunk.FromFileHandle.
                    ByteRange.Length.QuadPart =
                    HTTP_BYTE_RANGE_TO_EOF;

                dataChunk.FromFileHandle.FileHandle = hTempFile;

                result = HttpSendResponseEntityBody(
                    hReqQueue,
                    pRequest->RequestId,
                    0,           // This is the last send.
                    1,           // Entity Chunk Count.
                    &dataChunk,
                    NULL,
                    NULL,
                    0,
                    NULL,
                    NULL
                );

                if (result != NO_ERROR)
                {
                    wprintf(
                        L"HttpSendResponseEntityBody failed %lu\n",
                        result
                    );
                }

                goto Done;

                break;


            default:
                wprintf(
                    L"HttpReceiveRequestEntityBody failed with %lu \n",
                    result);
                goto Done;
            }

        } while (TRUE);
    }
    else
    {
        // This request does not have an entity body.
        //

        result = HttpSendHttpResponse(
            hReqQueue,           // ReqQueueHandle
            pRequest->RequestId, // Request ID
            0,
            &response,           // HTTP response
            NULL,                // pReserved1
            &bytesSent,          // bytes sent (optional)
            NULL,                // pReserved2
            0,                   // Reserved3
            NULL,                // LPOVERLAPPED
            NULL                 // pReserved4
        );
        if (result != NO_ERROR)
        {
            wprintf(L"HttpSendHttpResponse failed with %lu \n",
                result);
        }
    }

Done:

    if (pEntityBuffer)
    {
        FREE_MEM(pEntityBuffer);
    }

    if (INVALID_HANDLE_VALUE != hTempFile)
    {
        CloseHandle(hTempFile);
        DeleteFile(szTempName);
    }

    return result;
}

DWORD Win32WebServer::sendHttpResponse(PHTTP_REQUEST pRequest, USHORT StatusCode, const std::string& reason, const std::string& message)
{
    HTTP_RESPONSE response;
    RtlZeroMemory((&response), sizeof(*(&response)));
    response.StatusCode = StatusCode;
    response.pReason = reason.c_str();                         
    response.ReasonLength = (USHORT)strlen(reason.c_str());

    const std::string content_type = "text / html";
    response.Headers.KnownHeaders[HttpHeaderContentType].pRawValue = content_type.c_str();
    response.Headers.KnownHeaders[HttpHeaderContentType].RawValueLength = (USHORT)strlen(content_type.c_str());

    if (!message.empty())
    {
        auto entity_string = message;
        HTTP_DATA_CHUNK dataChunk;
        dataChunk.DataChunkType = HttpDataChunkFromMemory;
        dataChunk.FromMemory.pBuffer = entity_string.data();
        dataChunk.FromMemory.BufferLength = (ULONG)strlen(entity_string.data());

        response.EntityChunkCount = 1;
        response.pEntityChunks = &dataChunk;
    }

    DWORD bytesSent;
    DWORD result = ::HttpSendHttpResponse(
        mWorkingQueue,           // ReqQueueHandle
        pRequest->RequestId, // Request ID
        0,                   // Flags
        &response,           // HTTP response
        NULL,                // pReserved1
        &bytesSent,          // bytes sent  (OPTIONAL)
        NULL,                // pReserved2  (must be NULL)
        0,                   // Reserved3   (must be 0)
        NULL,                // LPOVERLAPPED(OPTIONAL)
        NULL                 // pReserved4  (must be NULL)
    );

    if (result != NO_ERROR)
    {
        MLOG_ERROR("Http response failed with error: " << result);
    }
    return result;
}

void Win32WebServer::run()
{
    // Allocate a 2 KB buffer. This size should work for most 
    // requests. The buffer size can be increased if required. Space
    // is also required for an HTTP_REQUEST structure.
    ULONG RequestBufferLength = sizeof(HTTP_REQUEST) + 2048;
    PCHAR pRequestBuffer = reinterpret_cast<PCHAR>(::HeapAlloc(::GetProcessHeap(), 0, (RequestBufferLength)));
    if (pRequestBuffer == nullptr)
    {
        MLOG_ERROR("Failed to allocate http request buffer");
        return;
    }
    PHTTP_REQUEST pRequest = (PHTTP_REQUEST)pRequestBuffer;

    // Wait for a new request. This is indicated by a NULL 
    // request ID.
    HTTP_REQUEST_ID    requestId;
    HTTP_SET_NULL_ID(&requestId);
    while(true)
    {
        RtlZeroMemory(pRequest, RequestBufferLength);

        DWORD              bytesRead;
        ULONG result = ::HttpReceiveHttpRequest(
            mWorkingQueue,          // Req Queue
            requestId,          // Req ID
            0,                  // Flags
            pRequest,           // HTTP request buffer
            RequestBufferLength,// req buffer length
            &bytesRead,         // bytes received
            nullptr                // LPOVERLAPPED
        );

        if (NO_ERROR == result)
        {
            switch (pRequest->Verb)
            {
            case HttpVerbGET:
                wprintf(L"Got a GET request for %ws \n",
                    pRequest->CookedUrl.pFullUrl);

                result = SendHttpResponse(
                    hReqQueue,
                    pRequest,
                    200,
                    "OK",
                    "Hey! You hit the server \r\n"
                );
                break;

            case HttpVerbPOST:

                wprintf(L"Got a POST request for %ws \n",
                    pRequest->CookedUrl.pFullUrl);

                result = SendHttpPostResponse(hReqQueue, pRequest);
                break;

            default:
                wprintf(L"Got a unknown request for %ws \n",
                    pRequest->CookedUrl.pFullUrl);

                result = SendHttpResponse(
                    hReqQueue,
                    pRequest,
                    503,
                    "Not Implemented",
                    NULL
                );
                break;
            }

            if (result != NO_ERROR)
            {
                break;
            }

            //
            // Reset the Request ID to handle the next request.
            //
            HTTP_SET_NULL_ID(&requestId);
        }
        else if (result == ERROR_MORE_DATA)
        {
            //
            // The input buffer was too small to hold the request
            // headers. Increase the buffer size and call the 
            // API again. 
            //
            // When calling the API again, handle the request
            // that failed by passing a RequestID.
            //
            // This RequestID is read from the old buffer.
            //
            requestId = pRequest->RequestId;

            //
            // Free the old buffer and allocate a new buffer.
            //
            RequestBufferLength = bytesRead;
            FREE_MEM(pRequestBuffer);
            pRequestBuffer = (PCHAR)ALLOC_MEM(RequestBufferLength);

            if (pRequestBuffer == NULL)
            {
                result = ERROR_NOT_ENOUGH_MEMORY;
                break;
            }

            pRequest = (PHTTP_REQUEST)pRequestBuffer;

        }
        else if (ERROR_CONNECTION_INVALID == result &&
            !HTTP_IS_NULL_ID(&requestId))
        {
            // The TCP connection was corrupted by the peer when
            // attempting to handle a request with more buffer. 
            // Continue to the next request.

            HTTP_SET_NULL_ID(&requestId);
        }
        else
        {
            break;
        }

    }

    if (pRequestBuffer)
    {
        FREE_MEM(pRequestBuffer);
    }
}