/////////////////////////////////////////////////////////////////////////////
//
//  Copyright (c) Microsoft Corporation.  All rights reserved.
//

#include "winlean.h"
#include "winsock2.h"
#include <stdlib.h>
#include <stdio.h>

#define DEF_PACKET_SIZE 32
#define MAX_PACKET 1024

#define xmalloc(s) HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, (s))
#define xfree(p)   HeapFree (GetProcessHeap(), 0, (p))

_declspec (thread) void * MC_MyCS_StopEvent;

typedef struct args {
    LPTSTR  host;
    int port;
} thread_args_t, *pthread_args_t;

static HANDLE  endThreadEvent;
static HANDLE  threadCountMutex;
static int doneThreads;
static int   numThreads;
static int  iterationCount;
static int  g_totalBytesRead;
/*

 * MC_GetLastErrorText
 *
 *  PURPOSE: copies error message text to string
 *
 *  PARAMETERS:
 *    lpszBuf - destination buffer
 *    dwSize - size of buffer
 *
 *  RETURN VALUE:
 *    destination buffer
 */
LPSTR MC_GetLastErrorText(LPSTR lpszBuf, DWORD dwSize, DWORD errnum)
{
    DWORD dwRet;
    LPTSTR lpszTemp = NULL;

    dwRet = FormatMessage( FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM |FORMAT_MESSAGE_ARGUMENT_ARRAY,
                           NULL,
                           errnum,
                           LANG_NEUTRAL,
                           (LPTSTR)&lpszTemp,
                           0,
                           NULL );

    if ( !dwRet)
        sprintf(lpszBuf, "error number %d", errnum); /* no message so use errnum */
    else
    {
        lpszTemp[lstrlen(lpszTemp)-2] = '\0';  /* remove cr and newline characters */
        //sprintf(lpszBuf, "%s (%d)", lpszTemp, errnum);
        sprintf(lpszBuf, " error %d\n", errnum);
    }

    if (lpszTemp)
        LocalFree((HLOCAL) lpszTemp );

    return lpszBuf;
}


static void MC_ErrMsg(LPSTR buf, DWORD bufsize, LPCSTR msg)
{
     DWORD tmp;

     tmp = strlen(strcpy(buf, msg));
     MC_GetLastErrorText(buf + tmp, bufsize - tmp, WSAGetLastError());
}

/*
 *  FUNCTION: MC_Recv
 *
 *  PURPOSE: A replacement for the Winsock recv() routine.  This version uses
 *           overlapped I/O to look for data on a socket and also watch
 *           for the Stop event.
 *
 *  PARAMETERS:
 *    sock - read from this socket.
 *    buffer - place incoming data into this buffer.
 *    nbytes - the amount of storage in the buffer.
 *
 *  RETURN VALUE:
 *    -1 - error.
 *    0  - the socket connection has been closed.
 *    >0 - the number of bytes read.
 */

int
MC_Recv(SOCKET sock, char *buffer, int nbytes)
{
    DWORD          bytesreceived;
    DWORD          flags = 0;
    WSAOVERLAPPED  ov;
    WSABUF         wsabuffer[1];
    HANDLE         ourevents[2];

    ov.hEvent = NULL;
    ov.Offset = ov.OffsetHigh = 0;
    wsabuffer[0].buf = buffer;
    wsabuffer[0].len = nbytes;

    while(1)
    {
        if (WSARecv(sock, wsabuffer, 1, &bytesreceived, &flags, &ov, NULL) == SOCKET_ERROR)
        {
            switch (WSAGetLastError())
            {
            case WSAEWOULDBLOCK: continue;  /* try again */
            case WSAENETRESET:
            case WSAECONNRESET:
            case WSAEDISCON:     return 0;  /* connection closed */
            //case WSA_IO_PENDING: break;     /* okay */
            default:             return -1; /* unexpected error */
            }
            ourevents[0] = (HANDLE) sock;
            ourevents[1] = MC_MyCS_StopEvent;
            if (WaitForMultipleObjects(2, ourevents, FALSE, 10000) != WAIT_OBJECT_0)
            {
                printf("MC_Recv(): Stop event or WaitForMultipleObjects() error\n");
                return -1; /* Stop event or error */
            }
            else if (!WSAGetOverlappedResult(sock, &ov, &bytesreceived, FALSE, &flags))
                return -1; /* error */
        }
        return bytesreceived; /* if we got here its due to a successful read */
    }
}

static DWORD  Gulp(SOCKET sockfd )
{
    int             totalBytes;
    int             numPackets;
    int             bytesRead;
    char            buffer[4096];

    totalBytes = 0;
    numPackets = 0;

    do {
        bytesRead = MC_Recv(sockfd, buffer,  4096);
        if (bytesRead > 0 ) {
            numPackets++;
            totalBytes += bytesRead;
        }
    } while (bytesRead > 0) ;

    //printf(" Read %d bytes in %d receives\n", totalBytes, numPackets);
    return totalBytes;
}


static DWORD WINAPI
MC_DoFetchUrl(LPVOID lpParam)
{
    struct hostent      *hostelement;
    struct sockaddr_in  sock_addr;
    DWORD               tmp;
    int                 server_status = 0;
    int                 ret_count = 0;
    SOCKET              sockfd;
    pthread_args_t      args;
    char                retmsg[1024];
    int                 size_retmsg;
    int                 i;
    struct linger       lo;
    int                 bytesRead;

    args = (pthread_args_t) lpParam;
    size_retmsg = 1024;

    bytesRead = 0;
    for (i=0; i < iterationCount; i++)  {

        memset((void *) &sock_addr, '\0', sizeof(sock_addr));
        sock_addr.sin_family = AF_INET;
        sock_addr.sin_port = htons((short) (args->port > 0 ? args->port : 80));

        if ((sockfd = WSASocket(AF_INET, SOCK_STREAM, IPPROTO_TCP, NULL, 0, 0)) == INVALID_SOCKET) {
            MC_ErrMsg(retmsg, size_retmsg, "500 Unable to create socket: ");
            printf(retmsg);
            //break;
        }

        if ((hostelement = gethostbyname(args->host)) == NULL) {
            tmp = sprintf(retmsg, "Unable to resolve host name '%s':", args->host);
            MC_GetLastErrorText(retmsg + tmp, size_retmsg - tmp, WSAGetLastError());
            printf(retmsg);
            break;
        }

        memcpy(&sock_addr.sin_addr, hostelement->h_addr, hostelement->h_length);
        tmp = connect(sockfd, (struct sockaddr *) &sock_addr, sizeof(sock_addr));
        if (tmp == SOCKET_ERROR) {
            MC_ErrMsg(retmsg, size_retmsg, "500 Unable to connect to web server: ");
            printf(retmsg);
            //break;
        }

        lo.l_onoff  = 1;
        lo.l_linger = 0;
        if (setsockopt(sockfd, SOL_SOCKET, SO_LINGER,
            (char *) &lo, sizeof(lo)) < 0) {
            printf("HTDoConnect: Can't set sockopt SO_LINGER (%d)\n", errno);
            break;
        }

        bytesRead += Gulp(sockfd);
        shutdown(sockfd, SD_BOTH);
        closesocket(sockfd);
    }

    WaitForSingleObject(threadCountMutex,INFINITE);
        g_totalBytesRead += bytesRead;
        doneThreads++;
        printf("Thread ending..count=%d\n",doneThreads);
        SetEvent(endThreadEvent);
    ReleaseMutex(threadCountMutex);

    return 1;
}

/* print an error message and exit the program */
void die (char *format,...) {
    va_list args;
    char buffer[2048], *ptr=buffer;

    va_start(args,format);
    ptr += vsprintf (ptr, format, args);
    va_end (args);
    *(ptr++) = '\n';
    *(ptr++) = '\0';

    printf( buffer);

    exit(1);
}
int InitChildLayer(void) {
    WORD wVersionRequested;
    WSADATA wsaData;
    int err;
    int major, minor;

    wVersionRequested = MAKEWORD( 2, 0 );

    err = WSAStartup( wVersionRequested, &wsaData );
    if ( err != 0 ) {
    /* Tell the user that we couldn't find a useable */
    /* WinSock DLL.                                  */
    die ("Couldn't find a WinSock DLL!\n");
    }

    /* Confirm that the WinSock DLL supports 2.0.        */
    /* Note that if the DLL supports versions greater    */
    /* than 2.0 in addition to 2.0, it will still return */
    /* 2.0 in wVersion since that is the version we      */
    /* requested.                                        */

    major = LOBYTE( wsaData.wVersion );
    minor = LOBYTE( wsaData.wVersion );

    if ( major < 1 || ( major == 1 && minor < 1 ) ) {
    /* Tell the user that we couldn't find a useable */
    /* WinSock DLL.                                  */
    WSACleanup( );
    die ("Require at least WinSock 1.1, your version is %d.%d!\n",
         major, minor);
    }
    return 0;
}

void Usage()
{
    printf(
            "Usage:\n"
            "    gulp [options] \n"
            "Options:\n"
            "    /p:N             -- Port to connect to (default=5000)\n"
            "    /a:nnn.nnn.nnn   -- address to connect to\n"
            "    /i:N             -- number of iterations\n"
            );
}
int main(int argc, char *argv[]) {

    int                 arg;
    pthread_args_t      blee;
    BOOL                fNeedHelp       = FALSE;
    char                *argn;
    char                *argp;
    char                *address;
    DWORD               dwPort;
    BOOL                bHaveAddress;
    LARGE_INTEGER       before;
    LARGE_INTEGER       after;
    LARGE_INTEGER       frequency;
    double              elapsed;
    double              rate;
    double              MBS;
    double              mbs;
    int                 i;
    DWORD               dwThreadId;
    HANDLE              hThread;

    InitChildLayer();

    endThreadEvent = CreateEvent(NULL,FALSE,FALSE,NULL);
    threadCountMutex = CreateMutex(NULL,FALSE,NULL);
    doneThreads = 0;

    if (endThreadEvent == NULL) {
        printf("Unable to allocate event!\n");
        exit(1);
    }

   numThreads = 1;
   dwPort = 5000;
   iterationCount = 10;
   fNeedHelp = FALSE;
   bHaveAddress = FALSE;

   for (arg = 1; arg < argc && !fNeedHelp; arg++) {
        if (argv[arg][0] == '-' || argv[arg][0] == '/') {
            argn = argv[arg]+1;                   // Argument name
            argp = argn;                          // Argument parameter

            while (*argp && *argp != ':') {
                argp++;
            }
            if (*argp == ':')
                *argp++ = '\0';

            switch (argn[0]) {

              case '?':                                 // Help
                fNeedHelp = TRUE;
                break;

              case 'a':                                 // address to connect to
              case 'A':
                address = argp;
                printf("Address = '%s'\n", argp);
                bHaveAddress = TRUE;
                break;

              case 'i':                                 // iterations within thread
              case 'I':
                iterationCount = atoi(argp);
                break;

              case 'p':                                 // port to attach to
              case 'P':
                dwPort = (DWORD)atoi(argp);
                break;

              case 't':                                 // Number of  threads to start
              case 'T':
                numThreads = atoi(argp);
                break;

              default:
                printf("Unknown argument: %s\n", argv[arg]);
                fNeedHelp = TRUE;
                break;
            }
        }
        else {
        }
    }

   if (fNeedHelp | !bHaveAddress) {
        Usage();
        exit(1);
   }

    printf ("Connecting to %s on port %d. threads=%d iterations=%d..", address, dwPort, numThreads, iterationCount);
    blee  = malloc(sizeof(pthread_args_t));
    blee->host = address;
    blee->port = dwPort;

#if 1
    for (i=0; i < numThreads; i++) {
        hThread = CreateThread(
            NULL,                       // default security attributes
            0,                          // use default stack size
            MC_DoFetchUrl,              // thread function
            blee,                       // argument to thread function
            0,                          // use default creation flags
            &dwThreadId);               // returns the thread identifier
    }
    //for (i=0; i < numThreads; i++) {
    QueryPerformanceFrequency(&frequency);
    QueryPerformanceCounter(&before);
    while(doneThreads !=  numThreads ) {
        WaitForSingleObject(endThreadEvent,INFINITE);
        printf ("received %d of %d events, done=%d\n", i+1, numThreads, doneThreads);
    }
    QueryPerformanceCounter(&after);
    after.QuadPart -= before.QuadPart;
    elapsed = (double)after.QuadPart / (double)frequency.QuadPart;
    rate = (double) g_totalBytesRead / elapsed;
    printf("bytes read =%d, time elapsed = %10.3f seconds rate = %10.3f \n", g_totalBytesRead, elapsed, rate);
    MBS =  rate / (double) (1024 *1024);
    mbs =  MBS * 8;
    printf("MB/s = %10.3f, mb/s = %10.3f\n", MBS, mbs);
#else
    MC_DoFetchUrl(blee);
#endif
    WSACleanup( );

}