#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <errno.h>
#include <arpa/inet.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <unistd.h>
#include <fcntl.h>
#include <poll.h>

#define ERRNO_EXIT \
    fprintf(stderr, "errno=%d (%s)\n", errno, strerror(errno)); \
    exit(EXIT_FAILURE)

int
main(int argc, const char *argv[])
{
    struct sockaddr_in6 addr;
    int                 udp_s;
    int                 tcp_s;
    int                 acc_s;
    int                 tcp_conn_s;
    int                 udp_port;
    int                 tcp_port;

    struct sockaddr_in6 bound_addr;
    socklen_t           addr_len;
    char                addr_str[200];
    char                port_str[200];

    struct pollfd       fds[2];
    int                 nfds;
    int                 rc;
    int                 tclass;

    if (argc < 3)
    {
        fprintf(stderr, "Usage: ./ipv6_class_server ipv6_addr tclass\n");
        exit(EXIT_FAILURE);
    }

    tclass = atoi(argv[2]);

    memset(&addr, sizeof(addr), 0);
    if (inet_pton(AF_INET6, argv[1], &addr.sin6_addr) < 1)
    {
        fprintf(stderr, "Failed to convert '%s' to IPv6 address\n",
                argv[1]);
        ERRNO_EXIT;
    }
    addr.sin6_family = AF_INET6;

    udp_s = socket(AF_INET6, SOCK_DGRAM, 0);
    if (udp_s < 0)
    {
        fprintf(stderr, "Failed to create UDP socket\n");
        ERRNO_EXIT;
    }

    tcp_s = socket(AF_INET6, SOCK_STREAM, 0);
    if (tcp_s < 0)
    {
        fprintf(stderr, "Failed to create TCP socket\n");
        ERRNO_EXIT;
    }

    if (setsockopt(udp_s, SOL_IPV6, IPV6_TCLASS, &tclass,
                   sizeof(tclass)) < 0)
    {
        fprintf(stderr, "setsockopt(IPV6_TCLASS) failed for UDP socket\n");
        ERRNO_EXIT;
    }

    if (setsockopt(tcp_s, SOL_IPV6, IPV6_TCLASS, &tclass,
                   sizeof(tclass)) < 0)
    {
        fprintf(stderr, "setsockopt(IPV6_TCLASS) failed for TCP socket\n");
        ERRNO_EXIT;
    }

    if (bind(udp_s, (struct sockaddr *)&addr, sizeof(addr)) < 0)
    {
        fprintf(stderr, "Failed to bind UDP socket\n");
        ERRNO_EXIT;
    }

    if (bind(tcp_s, (struct sockaddr *)&addr, sizeof(addr)) < 0)
    {
        fprintf(stderr, "Failed to bind TCP socket\n");
        ERRNO_EXIT;
    }

    addr_len = sizeof(bound_addr);
    if (getsockname(udp_s, (struct sockaddr *)&bound_addr,
                    &addr_len) < 0)
    {
        fprintf(stderr, "getsockname() failed for UDP socket\n");
        ERRNO_EXIT;
    }
    udp_port = ntohs(bound_addr.sin6_port);

    addr_len = sizeof(bound_addr);
    if (getsockname(tcp_s, (struct sockaddr *)&bound_addr,
                    &addr_len) < 0)
    {
        fprintf(stderr, "getsockname() failed for TCP socket\n");
        ERRNO_EXIT;
    }
    tcp_port = ntohs(bound_addr.sin6_port);

    printf("%s %d %d %d\n", argv[1], tclass, udp_port, tcp_port);

    if (listen(tcp_s, 1) < 0)
    {
        fprintf(stderr, "listen() failed for TCP socket\n");
        ERRNO_EXIT;
    }

    memset(fds, 0, sizeof(fds));
    fds[0].fd = tcp_s;
    fds[0].events = POLLIN;
    fds[1].fd = udp_s;
    fds[1].events = POLLIN;

    while (1)
    {
        rc = poll(fds, 2, -1);
        if (rc < 0)
        {
            fprintf(stderr, "poll() failed\n");
            ERRNO_EXIT;
        }

        /* TCP listener accepts connection */
        if (fds[0].revents == POLLIN)
        {
            addr_len = sizeof(addr);
            acc_s = accept(tcp_s, (struct sockaddr *)&addr,
                           &addr_len);
            if (acc_s < 0)
            {
                fprintf(stderr, "accept() failed\n");
                ERRNO_EXIT;
            }

            if (inet_ntop(AF_INET6, &addr.sin6_addr,
                          addr_str, sizeof(addr_str)) == NULL)
            {
                fprintf(stderr, "Failed to convert source address "
                        "of accepted connection to string\n");
                ERRNO_EXIT;
            }
            printf("Accepted connection from %s %d",
                   addr_str, (int)ntohs(addr.sin6_port));

            close(acc_s);
        }
        else if (fds[0].revents != 0)
        {
            fprintf(stderr, "Unexpected events 0x%x for TCP listener\n",
                    fds[0].revents);
            break;
        }

        /* UDP socket received data */
        if (fds[1].revents == POLLIN)
        {
            addr_len = sizeof(addr);
            rc = recvfrom(udp_s, port_str, sizeof(port_str), 0,
                          (struct sockaddr *)&addr, &addr_len);
            if (rc < 0)
            {
                fprintf(stderr, "recvfrom() failed\n");
                ERRNO_EXIT;
            }
            else if (rc == 0)
            {
                fprintf(stderr, "recvfrom() returned zero\n");
                exit(EXIT_FAILURE);
            }

            if (inet_ntop(AF_INET6, &addr.sin6_addr,
                          addr_str, sizeof(addr_str)) == NULL)
            {
                fprintf(stderr, "Failed to convert source address "
                        "of received data to string\n");
                ERRNO_EXIT;
            }

            printf("Received data from %s %d\n",
                   addr_str, (int)ntohs(addr.sin6_port));

            /*
             * Via UDP socket we receive port number to which
             * we try to establish TCP connection, so that
             * passive TCP connection establishment will be
             * checked by the client.
             */

            if (port_str[rc - 1] != '\0')
            {
                fprintf(stderr, "Not null-terminated string was received "
                        "from UDP socket\n");
                exit(EXIT_FAILURE);
            }
            printf("TCP socket will try to connect to %s %s\n",
                   addr_str, port_str);

            addr.sin6_port = htons(atoi(port_str));

            tcp_conn_s = socket(AF_INET6, SOCK_STREAM, 0);
            if (tcp_conn_s < 0)
            {
                fprintf(stderr, "Failed to create TCP socket\n");
                ERRNO_EXIT;
            }

            if (setsockopt(tcp_conn_s, SOL_IPV6, IPV6_TCLASS, &tclass,
                           sizeof(tclass)) < 0)
            {
                fprintf(stderr, "setsockopt(IPV6_TCLASS) failed for "
                        "TCP socket\n");
                ERRNO_EXIT;
            }

            bound_addr.sin6_port = 0;
            if (bind(tcp_conn_s, (struct sockaddr *)&bound_addr,
                     sizeof(bound_addr)) < 0)
            {
                fprintf(stderr, "Failed to bind TCP socket\n");
                ERRNO_EXIT;
            }

            if (fcntl(tcp_conn_s, F_SETFL, O_NONBLOCK) < 0)
            {
                fprintf(stderr, "Failed to make TCP socket non-blocking\n");
                ERRNO_EXIT;
            }

            rc = connect(tcp_conn_s, (struct sockaddr *)&addr,
                         sizeof(addr));
            if (rc < 0 && errno != EINPROGRESS)
            {
                fprintf(stderr, "Nonblocking connect() failed with "
                        "unexpected errno\n");
                ERRNO_EXIT;
            }

            sleep(1);
            rc = connect(tcp_conn_s, (struct sockaddr *)&addr,
                         sizeof(addr));
            if (rc < 0 && errno != EISCONN)
            {
                fprintf(stderr, "TCP connection was not established in "
                        "1 second\n");
            }

            close(tcp_conn_s);
        }
        else if (fds[1].revents != 0)
        {
            fprintf(stderr, "Unexpected events 0x%x for UDP socket\n",
                    fds[1].revents);
            break;
        }
    }

    close(tcp_s);
    close(tcp_conn_s);
    close(udp_s);

    return 0;
}