#include <windows.h>
#include <winsock.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define BUF_SIZE 4096

static SOCKET sock = INVALID_SOCKET;
static HANDLE in = INVALID_HANDLE_VALUE;
static DWORD old_mode;
static volatile LONG running = 1;

static BOOL WINAPI handler(DWORD type)
{
    if (type == CTRL_C_EVENT ||
        type == CTRL_BREAK_EVENT ||
        type == CTRL_CLOSE_EVENT ||
        type == CTRL_LOGOFF_EVENT ||
        type == CTRL_SHUTDOWN_EVENT)
    {
        InterlockedExchange(&running, 0);
        SetConsoleMode(in, old_mode);

        if (sock != INVALID_SOCKET) {
            closesocket(sock);
            sock = INVALID_SOCKET;
        }

        return TRUE;
    }

    return FALSE;
}

static DWORD WINAPI input_thread(LPVOID arg)
{
    char buf[256];
    DWORD n, sent;
    int r;

    (void)arg;

    while (running) {
        if (!ReadFile(in, buf, sizeof(buf), &n, NULL) || !n)
            break;

        sent = 0;

        while (sent < n && running) {
            r = send(sock, buf + sent, (int)(n - sent), 0);

            if (r <= 0) {
                InterlockedExchange(&running, 0);
                break;
            }

            sent += r;
        }
    }

    return 0;
}

int main(int argc, char **argv)
{
    WSADATA wsa;
    struct sockaddr_in addr;
    HANDLE thread;
    DWORD n;
    char buf[BUF_SIZE];
    int port, r;
    unsigned long ip;
    fd_set fds;
    struct timeval tv;

    if (argc != 3) {
        fprintf(stderr, "usage: %s <ip> <port>\n", argv[0]);
        return 1;
    }

    port = atoi(argv[2]);
    if (port < 1 || port > 65535) {
        fprintf(stderr, "invalid port\n");
        return 1;
    }

    if (WSAStartup(MAKEWORD(1, 1), &wsa) != 0 ||
        LOBYTE(wsa.wVersion) != 1) {
        fprintf(stderr, "Winsock 1.1 required\n");
        return 1;
    }

    in = GetStdHandle(STD_INPUT_HANDLE);

    if (!GetConsoleMode(in, &old_mode)) {
        fprintf(stderr, "not a console\n");
        WSACleanup();
        return 1;
    }

    SetConsoleMode(
        in,
        old_mode & ~(ENABLE_LINE_INPUT | ENABLE_ECHO_INPUT)
    );

    SetConsoleCtrlHandler(handler, TRUE);

    sock = socket(AF_INET, SOCK_STREAM, 0);
    if (sock == INVALID_SOCKET) {
        fprintf(stderr, "socket failed: %d\n", WSAGetLastError());
        SetConsoleMode(in, old_mode);
        WSACleanup();
        return 1;
    }

    ip = inet_addr(argv[1]);
    if (ip == INADDR_NONE) {
        fprintf(stderr, "invalid address\n");
        closesocket(sock);
        SetConsoleMode(in, old_mode);
        WSACleanup();
        return 1;
    }

    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = ip;
    addr.sin_port = htons((u_short)port);

    if (connect(sock, (struct sockaddr *)&addr, sizeof(addr)) == SOCKET_ERROR) {
        fprintf(stderr, "connect failed: %d\n", WSAGetLastError());
        closesocket(sock);
        SetConsoleMode(in, old_mode);
        WSACleanup();
        return 1;
    }

    thread = CreateThread(NULL, 0, input_thread, NULL, 0, NULL);

    if (!thread) {
        fprintf(stderr, "thread failed\n");
        closesocket(sock);
        SetConsoleMode(in, old_mode);
        WSACleanup();
        return 1;
    }

    while (running) {
        FD_ZERO(&fds);
        FD_SET(sock, &fds);

        tv.tv_sec = 0;
        tv.tv_usec = 100000;

        r = select(0, &fds, NULL, NULL, &tv);

        if (r == SOCKET_ERROR)
            break;

        if (!r)
            continue;

        r = recv(sock, buf, sizeof(buf), 0);

        if (r <= 0)
            break;

        WriteFile(
            GetStdHandle(STD_OUTPUT_HANDLE),
            buf,
            r,
            &n,
            NULL
        );
    }

    InterlockedExchange(&running, 0);

    closesocket(sock);
    sock = INVALID_SOCKET;

    WaitForSingleObject(thread, 250);
    CloseHandle(thread);

    SetConsoleMode(in, old_mode);
    SetConsoleCtrlHandler(handler, FALSE);
    WSACleanup();

    return 0;
}