#include <arpa/inet.h>
#include <fcntl.h>
#include <netdb.h>
#include <signal.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/epoll.h>
#include <sys/socket.h>
#include <termios.h>
#include <unistd.h>

#define SE 0xF0
#define SB 0xFA
#define WILL 0xFB
#define WONT 0xFC
#define DO 0xFD
#define DONT 0xFE
#define IAC 0xFF
#define CMD_WINDOW_SIZE 0x1F

int epoll_fd;
int sock;
struct termios tin;

void
cleanup(void)
{
    if (tcsetattr(STDIN_FILENO, TCSANOW, &tin))
        perror("ERROR: tcsetattr");

    if (close(epoll_fd) || close(sock))
        perror("ERROR: close");

    putchar('\r');
}

uint_fast8_t
negotiate(unsigned char* buf)
{
    if (buf[1] == DO && buf[2] == CMD_WINDOW_SIZE) {
        uint8_t msg0[3] = { IAC, WILL, CMD_WINDOW_SIZE };
        if (send(sock, msg0, 3, 0) != 3) {
            perror("ERROR: send");
            return 0;
        }

        uint8_t msg1[9] = { IAC, SB, CMD_WINDOW_SIZE, 0, 80, 0, 24, IAC, SE };
        if (send(sock, msg1, 9, 0) != 9) {
            perror("ERROR: send");
            return 0;
        }

        return 1;
    }

    if (buf[1] == DO)
        buf[1] = WONT;
    else if (buf[1] == WILL)
        buf[1] = DO;

    if (send(sock, buf, 3, 0) != 3) {
        perror("ERROR: send");
        return 0;
    }

    return 1;
}

uint_fast8_t
net_init(const char* host, const char* port)
{
    sock = socket(AF_INET, SOCK_STREAM, 0);
    if (sock == -1) {
        perror("ERROR: socket");
        return 0;
    }

    struct addrinfo hints;
    memset(&hints, 0, sizeof(struct addrinfo));
    hints.ai_family = AF_INET;
    hints.ai_socktype = SOCK_STREAM;

    struct addrinfo* res;
    int_fast8_t err = getaddrinfo(host, port, &hints, &res);
    if (err) {
        if (err == EAI_SYSTEM)
            perror("ERROR: getaddrinfo");
        else
            fprintf(stderr, "ERROR: getaddrinfo: %s\n", gai_strerror(err));
        return 0;
    }

    if (connect(sock, res->ai_addr, res->ai_addrlen)) {
        perror("ERROR: connect");
        if (close(sock))
            perror("ERROR: close");
        return 0;
    }

    freeaddrinfo(res);

    return 1;
}

void
sighandler(int sig)
{
    fputs("\n\rINFO: Received SIGTERM\n", stderr);
    exit(sig);
}

uint_fast8_t
terminal_set(void)
{
    if (tcgetattr(STDIN_FILENO, &tin)) {
        perror("ERROR: tcgetattr");
        return 0;
    }

    static struct termios tlocal;
    memcpy(&tlocal, &tin, sizeof(struct termios));
    cfmakeraw(&tlocal);

    if (tcsetattr(STDIN_FILENO, TCSANOW, &tlocal)) {
        perror("ERROR: tcsetattr");
        return 0;
    }

    return 1;
}

int
main(int argc, char** argv)
{
    if (argc < 2 || argc > 3) {
        printf("USAGE: %s, address|domainname [port]\n", argv[0]);
        return 1;
    }

    if (!net_init(argv[1], (argc == 3) ? argv[2] : "23"))
        return 2;

    if (!terminal_set())
        return 3;

    epoll_fd = epoll_create1(0);
    if (epoll_fd == -1) {
        perror("ERROR: epoll_create1");
        return 4;
    }

    atexit(cleanup);
    if (signal(SIGTERM, sighandler) == SIG_ERR) {
        perror("ERROR: signal");
        return 5;
    }

    struct epoll_event ev[2];
    memset(&ev, 0, 2 * sizeof(struct epoll_event));

    ev[0].events = EPOLLIN;
    ev[0].data.fd = 0;
    if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, 0, &ev[0])) {
        perror("ERROR: epoll_ctl");
        return 6;
    }

    ev[1].events = EPOLLIN;
    ev[1].data.fd = sock;
    if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sock, &ev[1])) {
        perror("ERROR: epoll_ctl");
        return 6;
    }

    uint_fast8_t run = 1;

    do {
        struct epoll_event events[2];
        int_fast8_t count = epoll_wait(epoll_fd, events, 2, -1);
        if (count == -1) {
            perror("ERROR: epoll_wait");
            return 7;
        }

        for (int_fast8_t i = 0; i < count; ++i) {
            if (events[i].data.fd == sock) {
                uint8_t buf[3];
                memset(buf, 0, 3);
                int_fast8_t ret = recv(sock, buf, 1, 0);
                if (ret == -1) {
                    perror("ERROR: recv");
                    return 8;
                }

                if (!ret)
                    return 0;

                if (buf[0] == IAC) {
                    ret = recv(sock, buf + 1, 2, 0);
                    if (ret < 0) {
                        perror("ERROR: recv");
                        return 8;
                    }

                    if (!ret)
                        return 0;

                    negotiate(buf);
                } else {
                    putchar(buf[0]);
                    fflush(stdin);
                }
            } else if (events[i].data.fd == 0) {
                const uint8_t chr = getchar();
                if (chr == 3) {
                    fputs("\n\rLeaving...\n\r", stdout);
                    return 0;
                }

                if (send(sock, &chr, 1, 0) != 1) {
                    perror("ERROR: send");
                    return 9;
                }

                if (chr == '\n')
                    putchar('\r');
            }
        }
    } while (run);
}