//
// A simple file server application.
// It listens to the port written in command line (default 1234),
// accepts a connection, and perform clients commands
//
// Usage:
//      server [port_to_listen]
// Default is the port 1234.
//
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <errno.h>
#include <assert.h>
#include <direct.h>
#include <limits.h>
#include <winsock2.h>
#include <ws2tcpip.h>
#include "filesrv.h"

static void usage();
static DWORD WINAPI serveClient(void* pSocket);
static void initializeFrame(ProtocolFrame* frame);

class ProtocolException {
public:
    const char *reason;
    ProtocolException():
        reason("")
    {}

    ProtocolException(const char *cause):
        reason(cause)
    {}
};

static int receiveFrame(int s, ProtocolFrame* frame);
static void sendFrame(int s, const ProtocolFrame* frame)
    throw (ProtocolException);

// Functions performing client commands:
static void receiveFile(int s,  ProtocolFrame* frame) throw (ProtocolException);
static void sendFile(int s, ProtocolFrame* frame) throw (ProtocolException);
static void pwd(int s, ProtocolFrame* frame) throw (ProtocolException);
static void cd(int s, ProtocolFrame* frame) throw (ProtocolException);
static void ls(int s, ProtocolFrame* frame) throw (ProtocolException);
static void protocolError(int s, const char *txt) throw (ProtocolException);
static void replySuccess(int s) throw (ProtocolException);

static int ascii2UTF16(
    const char * asciiStr,
    wchar_t* utf16Str,
    int maxLen
);

static const int BLOCK_SIZE = 1024;

int main(int argc, char *argv[]) {
    if (argc > 1 && *(argv[1]) == '-') {
        usage(); exit(1);
    }

    // Initialize Winsock DLL
    WORD wVersionRequested = MAKEWORD(2, 2);
    WSADATA wsaData;
    int err = WSAStartup(wVersionRequested, &wsaData);
    if (err != 0) {
        printf("WSAStartup failed with error: %d\n", err);
        exit(-1);
    }

    int listenPort = 1234;
    if (argc > 1)
        listenPort = atoi(argv[1]);

    // Create a socket
    int s0 = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (s0 < 0) {
        perror("Cannot create a socket"); exit(1);
    }

    // Fill in the address structure containing self address
    struct sockaddr_in myaddr;
    memset(&myaddr, 0, sizeof(struct sockaddr_in));
    myaddr.sin_family = AF_INET;
    myaddr.sin_port = htons(listenPort);        // Port to listen
    myaddr.sin_addr.s_addr = htonl(INADDR_ANY);

    // Bind a socket to the address
    int res = bind(s0, (struct sockaddr*) &myaddr, sizeof(myaddr));
    if (res < 0) {
        printf(
            "Cannot bind a socket: err=%d\n",
            WSAGetLastError()
        );
        exit(-1);
    }

    // Set the "LINGER" timeout to zero, to close the listen socket
    // immediately at program termination.
    struct linger linger_opt = { 1, 0 }; // Linger active, timeout 0
    setsockopt(
        s0, SOL_SOCKET, SO_LINGER,
        (const char*) &linger_opt, sizeof(linger_opt)
    );

    printf("File server - at your service.\n");
    fflush(stdout);

    // Now, listen for a connection
    res = listen(s0, 5);    // "5" is the maximal length of the queue
    if (res < 0) {
        perror("Cannot listen"); exit(1);
    }

    // Accept a connection (the "accept" command waits for a connection with
    // no timeout limit...)
    while (true) {
        struct sockaddr_in peeraddr;
        int peeraddr_len = sizeof(peeraddr);
        int s1 = accept(s0, (struct sockaddr*) &peeraddr, &peeraddr_len);
        if (s1 < 0) {
            printf(
                "Cannot accept: err=%d\n",
                WSAGetLastError()
            );
            exit(-1);
        }


        // A connection is accepted. The new socket "s1" is created
        // for data input/output. The peeraddr structure is filled in with
        // the address of connected entity, print it.

        int* pSocket = new int(s1);
        HANDLE hClientThread = NULL;
        DWORD clientThreadID = 0;
        hClientThread = CreateThread(
            NULL,             // lpThreadAttributes
            0,                // dwStackSize - default
            &serveClient,     // thread starting function
            (void *) pSocket, // parameter of thread starting function
            0,                // dwCreationFlags
            &clientThreadID
        );
        if (hClientThread == NULL) {
            printf(
                "Cannot create client thread: error %d\n",
                (int) GetLastError()
            );
            exit(-1);
        }
        CloseHandle(hClientThread);

        printf(
            "Connection from IP %d.%d.%d.%d, port %d\n",
            (int)(ntohl(peeraddr.sin_addr.s_addr) >> 24) & 0xff, // High byte of address
            (int)(ntohl(peeraddr.sin_addr.s_addr) >> 16) & 0xff, // . . .
            (int)(ntohl(peeraddr.sin_addr.s_addr) >> 8) & 0xff,  // . . .
            (int)ntohl(peeraddr.sin_addr.s_addr) & 0xff,         // Low byte of addr
            (int)ntohs(peeraddr.sin_port)
        );
        fflush(stdout);

        // Continue listening...
    }

    closesocket(s0); // Close the listen socket

    WSACleanup();
    return 0;
}

static void usage() {
    printf(
        "A test Internet file server.\n"
        "Usage:\n"
        "     server [port_to_listen]\n"
        "Default is the port 1234.\n"
    );
}

static DWORD WINAPI serveClient(void* pSocket) {
    int s = *((int *) pSocket);
    delete (int *) pSocket;

    ProtocolFrame frame;
    try {
        initializeFrame(&frame);
        frame.command = FSR_READY;
        frame.length = 0;
        sendFrame(s, &frame);

        while (receiveFrame(s, &frame) >= 0) {
            // printf("Command received=%d\n", frame.command);
            // fflush(stdout);

            if (frame.command == FSR_LOGOUT) {
                printf("Logout.\n");
                fflush(stdout);
                break;
            }

            switch (frame.command) {
            case FSR_PUT:
                receiveFile(s, &frame);
                break;
            case FSR_GET:
                sendFile(s, &frame);
                break;
            case FSR_PWD:
                pwd(s, &frame);
                break;
            case FSR_CD:
                cd(s, &frame);
                break;
            case FSR_LS:
                ls(s, &frame);
                break;
            default:
                protocolError(s, "Illegal command");
                break;
            }

            initializeFrame(&frame);
            frame.command = FSR_READY;
            frame.length = 0;
            sendFrame(s, &frame);
        }
    } catch (ProtocolException& e) {
        printf("Protocol exception...\n");
        printf("%s\n", e.reason);
        fflush(stdout);
    }

    printf("Closing a connection.\n");

    closesocket(s); // Close a connection
    return 0;
}

static void initializeFrame(ProtocolFrame* frame) {
    memset(frame, 0, sizeof(ProtocolFrame));
}

static void receiveFile(int s,  ProtocolFrame* frame)
    throw (ProtocolException) {
    wprintf(L"put %s\n", frame->data);
    fflush(stdout);

    wchar_t path[PATH_MAX+2];
    int len = frame->length/2;
    if (len > PATH_MAX)
        len = PATH_MAX;
    memmove(path, frame->data, len*2);
    path[len] = 0;
    FILE* f = _wfopen(path, L"wb");
    if (f == NULL) {
        perror("Cannot open a file for writing");
        wprintf(L"Path: %s\n", path);
        fflush(stdout);
        protocolError(s, strerror(errno));
        return;
    } else {
        initializeFrame(frame);
        frame->command = FSR_READY;
        frame->length = 0;
        sendFrame(s, frame);
    }

    int errnum = 0;
    while (true) {
        int res = receiveFrame(s, frame);
        if (res < 0)
            break;
        if (frame->command != FSR_DATA && frame->command != FSR_DATAEND)
            break;
        if (f != NULL) {
            if (fwrite(frame->data, 1, frame->length, f) <= 0) {
                errnum = errno;
                perror("Write error");
                fclose(f); f = NULL;
            }
        }
        if (frame->command == FSR_DATAEND)
            break;
    }

    if (f != NULL) {
        fclose(f);
        replySuccess(s);
    } else {
        protocolError(s, strerror(errnum));
    }
}

static void sendFile(int s, ProtocolFrame* frame)
    throw (ProtocolException) {
    wprintf(L"get %s\n", frame->data);
    fflush(stdout);

    wchar_t path[PATH_MAX+2];
    int len = frame->length/2;
    if (len > PATH_MAX)
        len = PATH_MAX;
    memmove(path, frame->data, 2*len);
    path[len] = 0;
    FILE* f = _wfopen(path, L"rb");
    if (f == NULL) {
        protocolError(s, strerror(errno));
        return;
    }
    bool endFile = false;
    while (!endFile) {
        initializeFrame(frame);
        int res = fread(frame->data, 1, BLOCK_SIZE, f);
        if (ferror(f)) {
            protocolError(s, strerror(errno));
            return;
        }
        if (feof(f)) {
            frame->command = FSR_DATAEND;
            endFile = true;
        } else {
            frame->command = FSR_DATA;
        }
        frame->length = res;
        sendFrame(s, frame);
    }
    fclose(f);
    //... replySuccess(s);
}

static void pwd(int s, ProtocolFrame* frame)
    throw (ProtocolException) {
    printf("pwd\n");
    fflush(stdout);

    wchar_t path[PATH_MAX+2];
    initializeFrame(frame);
    if (GetCurrentDirectoryW(PATH_MAX, path) == 0) {
        protocolError(s, strerror(errno));
        return;
    }
    int len = wcslen(path);
    if (len > FSR_MAXDATALEN/2 - 1) {
        len = FSR_MAXDATALEN/2 - 1;
    }
    path[len] = 0;
    frame->command = FSR_DATAEND;
    frame->length = 2*(len + 1);
    memmove(frame->data, path, 2*(len + 1));
    sendFrame(s, frame);
}

static void cd(int s, ProtocolFrame* frame)
    throw (ProtocolException) {

    wprintf(L"cd %s\n", frame->data);
    /*...
    wprintf(
        L"frame.length=%d data=%s\n",
        frame->length, frame->data
    );
    ...*/
    fflush(stdout);

    wchar_t path[PATH_MAX+2];
    assert(frame->command == FSR_CD);
    int len = frame->length/2;
    if (len > PATH_MAX)
        len = PATH_MAX;
    memmove(path, frame->data, 2*len);
    path[len] = 0;

    int res = _wchdir(path);
    if (res == 0) {
        replySuccess(s);
    } else {
        protocolError(s, strerror(errno));
    }
}

static void ls(int s, ProtocolFrame* frame)
    throw (ProtocolException) {
    printf("ls\n");
    fflush(stdout);

    char errorText[128];
    WIN32_FIND_DATAW findData;
    HANDLE h;
    h = FindFirstFileW(L"*", &findData);
    if (h == INVALID_HANDLE_VALUE) {
        DWORD err = GetLastError();
        DWORD res = FormatMessageA(
            FORMAT_MESSAGE_FROM_SYSTEM |
                FORMAT_MESSAGE_MAX_WIDTH_MASK, // flags
            NULL,                   // lpSource
            err,                    // dwMessageId,
            0,                      // dwLanguageId,
            errorText,
            126,                    // nSize,
            NULL                    // va_list *Arguments
        );
        if (res == 0) {
            sprintf(
                errorText,
                "Cannot list the directory: error %d",
                (int) err
            );
        }

        protocolError(s, errorText);
        return;
    }

    while (h != INVALID_HANDLE_VALUE) {
        initializeFrame(frame);
        memmove(frame->data, &findData.dwFileAttributes, 4);
        int nameLen = wcslen(findData.cFileName);
        memmove(frame->data + 4, findData.cFileName, 2*(nameLen + 1));
        frame->length = 4 + 2*(nameLen+1);
        frame->data[4 + 2*nameLen] = 0;
        frame->data[4 + 2*nameLen + 1] = 0;

        bool notLast = FindNextFileW(h, &findData);
        if (notLast)
            frame->command = FSR_DATA;
        else
            frame->command = FSR_DATAEND;
        sendFrame(s, frame);
        if (!notLast)
            break;
    }
    FindClose(h);
}

static void protocolError(int s, const char *txt)
    throw (ProtocolException) {
    ProtocolFrame frame;
    initializeFrame(&frame);
    frame.command = FSR_ERROR;
    const char *tx = txt;
    char errTxt[64];
    if (tx == NULL) {
        memset(errTxt, 0, 64);
        sprintf(errTxt, "Unknown error %d", s);
        tx = errTxt;
    }

    int len = ascii2UTF16(
        tx, (wchar_t *) frame.data, FSR_MAXDATALEN/2 - 1
    );
    frame.length = 2*(len + 1);
    sendFrame(s, &frame);
}

static void replySuccess(int s)  throw (ProtocolException) {
    ProtocolFrame frame;
    initializeFrame(&frame);
    frame.command = FSR_SUCCESS;
    frame.length = 0;
    sendFrame(s, &frame);
}

static void sendFrame(int s, const ProtocolFrame* frame)
    throw (ProtocolException) {
    ssize_t res = send(
        s, (char *)frame, FSR_HEADERLEN + frame->length, 0
    );
    if (res < 0) {
        printf(
            "Cannot send data: err=%d\n",
            WSAGetLastError()
        );
        fflush(stdout);
        throw ProtocolException("Send error on socket");
    }
}

static int receiveFrame(int s, ProtocolFrame* frame) {
    // Read frame header
    //???!!! ssize_t res = recv(s, frame, FSR_HEADERLEN, MSG_WAITALL);
    //??? #define MSG_WAITALL 0x8
    ssize_t res = 0;
    int lh = 0;
    do {
        res = recv(
            s,
            ((char *)frame) + lh,
            FSR_HEADERLEN - lh,
            0
        );
        if (res >= 0)
            lh += res;
    } while (res >= 0 && lh < FSR_HEADERLEN);

    /*
    printf(
        "Frame header received: len=%d command=%d\n",
        res, frame->command
    );
    fflush(stdout);
    */

    if (lh < FSR_HEADERLEN) {
        printf(
            "Could not receive a frame header: len=%d\n",
            (int) res
        );
        fflush(stdout);
        return (-1);
    }
    int len = frame->length;
    if (len > FSR_MAXDATALEN)
        len = FSR_MAXDATALEN;
    char *pos = frame->data;
    while (len > 0) {
        int bytesToRead = len;
        if (bytesToRead > BLOCK_SIZE)
            bytesToRead = BLOCK_SIZE;
        res = recv(s, pos, bytesToRead, 0);
        if (res <= 0)
            return (-1);
        len -= res;
        pos += res;
    }
    return (FSR_HEADERLEN + frame->length);
}

static int ascii2UTF16(
    const char * asciiStr,
    wchar_t* utf16Str,
    int maxLen
) {
    int len = strlen(asciiStr);
    if (len > maxLen)
        len = maxLen;
    for (int i = 0; i < len; ++i) {
        utf16Str[i] = (((wchar_t) asciiStr[i]) & 0xff);
    }
    utf16Str[len] = 0;
    return len;
}