// Troy Butler & Owen Sullivan
// CPSC 3600-001 F23
// Assignment 2

#include <float.h>
#include <getopt.h>
#include <math.h>
#include <netdb.h>
#include <pthread.h>
#include <signal.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <unistd.h>

// Define MAX_DATAGRAM_SIZE as 65507 bytes per maximum size of UDP data length
// Source: https://en.wikipedia.org/wiki/User_Datagram_Protocol
#define MAX_DATAGRAM_SIZE 65507

// Fix issue with getopt_long() not parsing optargs for optional arguments
// when whitespace exists between argument and optarg
// Source: https://cfengine.com/blog/2021/optional-arguments-with-getopt-long/
#define OPTARG_PRESENT                                           \
    ((optarg == NULL && optind < argc && argv[optind][0] != '-') \
         ? (bool)(optarg = argv[optind++])                       \
         : (optarg != NULL))

// Struct to store values of tracking variables for whether command-line 
// options were passed
struct optTrack {

    bool p;
    bool c;
    bool i;
    bool s;

} optTrack_default = {false, false, false, false};
// Typedef the struct as default
typedef struct optTrack optTrack_t;

// Struct to store the values of command-line arguments that have been passed
struct optStore {
    
    // Variables that store the status of non-valued arguments
    bool S;
    bool n;

    // Variables that store optargs as they are received
    char *p;
    unsigned int c;
    double i;
    size_t s;
    char *addr;

} optStore_default = {false, false, "33333", 0x7fffffff, 1.0, 12, NULL};
// Typedef the struct as default
typedef struct optStore optStore_t;

// Create variable for socket descriptors (created before calls to
// signal() so that the sockets can be closed on Ctrl + C)
int serverSocket;
int clientSocket;

// Struct to store values that need to be passed to thread functions
typedef struct threadArgs {
    
    optStore_t *optValues;
    int *socket;
    struct addrinfo *serverAddress;
    struct sockaddr_storage *clientAddress;
    socklen_t *clientAddressLength;

} threadArgs_t;

// Enumeration to make tracking active thread easier
enum thread { SENDER, RECEIVER };
enum thread currentThread = SENDER;

// Create a time representation
struct timespec start;

// Variables needed to ensure proper thread synchronization
pthread_mutex_t mutexLock = PTHREAD_MUTEX_INITIALIZER;
pthread_cond_t cond = PTHREAD_COND_INITIALIZER;

// Struct to keep track of ping statistics
typedef struct ping {
    
    struct timespec sentTime;
    bool received;
    struct timespec receivedTime;
    double difference;

} ping_t;

// Variable to count how many pings are sent
int numPings = 0;
// Variable to store all ping statistics
ping_t **pings;

// Function prototypes
void printErrorAndExit(char *message);
char *getOptarg(int argc, char **argv);
void printUsageAndExit();
void checkOptions(int argc, char **argv, optStore_t *optValues,
                  optTrack_t *optTracker);
void printOptions(optStore_t *optValues, optTrack_t *optTracker);
double calculateTimeDifference(struct timespec *start, struct timespec *end);
void *clientSender(void *threadArg);
void *clientReceiver(void *threadArg);
void exitClient();
void *serverReceiver(void *threadArg);
void *serverSender(void *threadArg);
void exitServer();

// Main function
int main(int argc, char **argv) {
    
    // Create an object to track the use of command-line arguments
    optTrack_t optTracker = optTrack_default;
    // Create an object to store the values of command-line arguments
    optStore_t optValues = optStore_default;
    // Call helper function that parses command-line arguments
    checkOptions(argc, argv, &optValues, &optTracker);

    // Tell the system what kind of server address we're expecting
    struct addrinfo addressCriteria; // Criteria for address match
    memset(&addressCriteria, 0, sizeof(addressCriteria)); // Zero out structure
    addressCriteria.ai_family = AF_UNSPEC; // Any address family
    addressCriteria.ai_socktype = SOCK_DGRAM; // Only datagram sockets
    addressCriteria.ai_protocol = IPPROTO_UDP; // Only UDP protocol

    // Create an object representing a list of addresses
    struct addrinfo *serverAddress;

    // Running in server mode
    if (optValues.S) {
        
        // Handle Ctrl + C and error check
        if (signal(SIGINT, exitServer) == SIG_ERR) {
            
            printErrorAndExit("ERROR: Cannot catch SIGINT (Ctrl + C)");
            
        }

        // Add extra address criteria flags for server
        addressCriteria.ai_flags = AI_PASSIVE;  // Accept on any address/port

        // Get server address info and error check
        if (getaddrinfo(NULL, optValues.p, &addressCriteria, &serverAddress) !=
            0) {
                
            printErrorAndExit("getaddrinfo() failed");
            
        }

        // Create socket for incoming connections
        serverSocket =
            socket(serverAddress->ai_family, serverAddress->ai_socktype,
                   serverAddress->ai_protocol);
        // Error check socket()
        if (serverSocket < 0) {
            
            printErrorAndExit("socket() failed");
            
        }

        // Bind to the local address and error check
        if (bind(serverSocket, serverAddress->ai_addr,
                 serverAddress->ai_addrlen) < 0) {
                     
            printErrorAndExit("bind() failed");
            
        }

        // Free address list allocated by getaddrinfo()
        freeaddrinfo(serverAddress);

        // Create an object representing the client address
        struct sockaddr_storage clientAddress;
        // Store the length of the client address information (somehow
        // required for recvfrom() to work)
        socklen_t clientAddressLength = sizeof(clientAddress);

        // Create an object to store values that need to be passed to thread
        // functions
        threadArgs_t threadArgs = {&optValues, &clientSocket, NULL,
                                   &clientAddress, &clientAddressLength};

        // Create threads for server receiving and sending
        pthread_t receiver;
        pthread_t sender;

        // Loop indefinitely to continue receiving packets from potential
        // clients
        while (true) {
            
            // Set the current thread to the receiver thread (server always
            // receives first)
            currentThread = RECEIVER;

            // Create thread contexts for server receiving and sending, passing
            // thread functions with thread arguments
            pthread_create(&receiver, NULL, serverReceiver, &threadArgs);
            pthread_create(&sender, NULL, serverSender, &threadArgs);

            // Switch contexts to server receiving and sending threads
            pthread_join(receiver, NULL);
            pthread_join(sender, NULL);
            
        }

        // Clean up thread mutex lock and condition variables
        pthread_mutex_destroy(&mutexLock);
        pthread_cond_destroy(&cond);

        // Running in client mode
    } else {
        
        // Handle Ctrl + C and error check
        if (signal(SIGINT, exitClient) == SIG_ERR) {
            
            printErrorAndExit("ERROR: Cannot catch SIGINT (Ctrl + C)");
            
        }

        // Store the system's current clock time
        clock_gettime(CLOCK_REALTIME, &start);

        // Call helper function that prints command-line argument values
        printOptions(&optValues, &optTracker);

        // Get server address info and error check
        if (getaddrinfo(optValues.addr, optValues.p, &addressCriteria,
                        &serverAddress) != 0) {
                            
            printErrorAndExit("getaddrinfo() failed");
            
        }

        // Create a socket for outgoing connections
        clientSocket =
            socket(serverAddress->ai_family, serverAddress->ai_socktype,
                   serverAddress->ai_protocol);
        // Error check socket()
        if (clientSocket < 0) {
            
            printErrorAndExit("socket() failed");
            
        }

        // Allocate memory for ping statistic storage based on number of pings
        // to be sent
        pings = malloc(optValues.c * sizeof(ping_t *));

        // Create an object to store values that need to be passed to thread
        // functions
        threadArgs_t threadArgs = {&optValues, &clientSocket, serverAddress,
                                   NULL, NULL};

        // Create threads for client sending and receiving
        pthread_t sender;
        pthread_t receiver;

        // Create thread contexts for client sending and receiving, passing
        // thread functions with thread arguments
        pthread_create(&sender, NULL, clientSender, &threadArgs);
        pthread_create(&receiver, NULL, clientReceiver, &threadArgs);

        // Switch contexts to client sending and receiving threads
        pthread_join(sender, NULL);
        pthread_join(receiver, NULL);

        // Clean up thread mutex lock and condition variables
        pthread_mutex_destroy(&mutexLock);
        pthread_cond_destroy(&cond);
        
    }
    
}

// Function that prints an error message and exits the program early
void printErrorAndExit(char *message) {
    
    fprintf(stderr, "%s\n", message);

    exit(1);
    
}

// Function to make the process of checking for an optional optarg with
// whitespace easier
char *getOptarg(int argc, char **argv) {
    
    return OPTARG_PRESENT ? optarg : NULL;
    
}

// Function that prints the program's intended command-line usage and exits
// early
void printUsageAndExit() {
    
    fprintf(stderr, "\n");
    fprintf(stderr, "Usage:\n");
    fprintf(stderr, "\tServer:\tudping -S -p <port_number> [-n]\n");
    fprintf(stderr,
            "\tClient:\tudping -p <port_number> -c <ping_packet_count> -i "
            "<ping_interval> -s <size_in_bytes> [-n] <server_ip_address>\n");

    exit(1);
    
}

// Function that checks for the correct required command-line arguments
void checkOptions(int argc, char **argv, optStore_t *optValues,
                  optTrack_t *optTracker) {
    
    // Variables for tracking with getopt_long()
    int option;
    int optionIndex;

    // Array of command-line options that can be parsed by getopt_long()
    static struct option options[] = {

        {"port-number", required_argument, NULL, 'p'},
        {"ping-packet-count", required_argument, NULL, 'c'},
        {"ping-interval", required_argument, NULL, 'i'},
        {"size-in-bytes", required_argument, NULL, 's'},
        {"no-print", no_argument, NULL, 'n'},
        {"server", no_argument, NULL, 'S'},
        {0, 0, 0, 0}

    };

    // Read command-line arguments while there is a valid argument to be parsed
    while ((option = getopt_long(argc, argv, "p:c::i::s::nS", options,
                                 &optionIndex)) != -1) {
                                     
        // Handle command-line arguments
        switch (option) {
            
            // Argument -p or --port-number
            case 'p': {
                
                optTracker->p = true;
                optValues->p = getOptarg(argc, argv);

            } break;

            // Argument -c or --ping-packet-count
            case 'c': {
                
                optTracker->c = true;
                optValues->c = atoi(getOptarg(argc, argv));

            } break;

            // Argument -i or --ping-interval
            case 'i': {
                
                optTracker->i = true;
                optValues->i = atof(getOptarg(argc, argv));

            } break;

            // Argument -s or --size-in-bytes
            case 's': {
                
                optTracker->s = true;
                sscanf(getOptarg(argc, argv), "%zu", &optValues->s);

            } break;

            // Argument -n or --no-print
            case 'n': {
                
                optValues->n = true;

            } break;

            // Argument -S or --server
            case 'S': {
                
                optValues->S = true;

            } break;

            // Default response (handles non-recognized options)
            default: {
                
                printUsageAndExit();

            } break;
            
        }
        
    }

    // If the -S or --server argument was not passed
    if (!optValues->S) {
        
        // If the index of the first non-opt command-line argument is less than
        // the total number of arguments, indicating that a value for the
        // server IP address was provided
        if (optind < argc) {
            
            // Store last command-line argument (not handled by getopt_long())
            // as the IP address for the server
            optValues->addr = argv[argc - 1];

            // Otherwise, a server IP address was not provided
        } else {
            
            fprintf(stderr, "ERROR: Missing server IP address\n");
            printUsageAndExit();
            
        }
        
    }
    
}

// Function that echoes selected command-line argument values to stderr
void printOptions(optStore_t *optValues, optTrack_t *optTracker) {
    
    // If the count was provided, print it
    if (optTracker->c) {

        fprintf(stderr, "%-9s %15d\n", "Count", optValues->c);

    }

    // If the size was provided, print it
    if (optTracker->s) {

        fprintf(stderr, "%-9s %15zu\n", "Size", optValues->s);

    }

    // If the interval was provided, print it
    if (optTracker->i) {

        fprintf(stderr, "%-9s %15.3lf\n", "Interval", optValues->i);

    }

    // If the port was provided, print it
    if (optTracker->p) {

        fprintf(stderr, "%-9s %15s\n", "Port", optValues->p);

    }

    // If the no-print was not provided, print the server address
    if (!optValues->n) {

        fprintf(stderr, "%-9s %15s\n", "Server_ip", optValues->addr);

        // If the no-print option was provided
    } else {

        // If any printable command-line options were provided, print a new 
        // line for spacing
        if (optTracker->c || optTracker->s || optTracker->i || optTracker->p) {

            fprintf(stdout, "\n");

        }
        
        // Print ten asterisks
        fprintf(stdout, "**********\n");

    }

    fprintf(stderr, "\n");
    
}

// Function that calculates the difference between two timespec objects
double calculateTimeDifference(struct timespec *start, struct timespec *end) {
    
    // Create a temporary time representation to hold difference values
    struct timespec difference;

    // If subtracting the nanoseconds would result in a negative value
    if ((end->tv_nsec - start->tv_nsec) < 0) {
        
        // Set the seconds equal to the difference of seconds minus 1
        difference.tv_sec = (end->tv_sec - start->tv_sec) - 1;
        // Set the nanoseconds equal to 1 second minus the difference of
        // nanoseconds
        difference.tv_nsec = 1e+9 + (end->tv_nsec - start->tv_nsec);

        // Otherwise, simply compute the differences for both seconds and
        // nanoseconds
    } else {
        
        difference.tv_sec = end->tv_sec - start->tv_sec;
        difference.tv_nsec = end->tv_nsec - start->tv_nsec;
        
    }

    // Return the equivalent value of seconds + 0.xxxx nanoseconds
    return (difference.tv_sec + (difference.tv_nsec / 1e+9));
    
}

// Function that performs client ping sending logic
void *clientSender(void *threadArg) {
    
    // Convert and store thread argument as struct object of thread argument
    // values
    threadArgs_t *threadArgs = threadArg;

    // Loop through the number of pings to be sent
    for (int i = 1; i <= threadArgs->optValues->c; i++) {
        
        // Lock mutex lock for global variable safety
        pthread_mutex_lock(&mutexLock);

        // Create another time representation and set it equal to the starting
        // value
        struct timespec interval = start;

        // Loop while the sender thread is not the active thread
        while (currentThread != SENDER) {
            
            // Calculate the offset to wait to send the next packet
            double offset = (i - 1) * threadArgs->optValues->i;
            // If the offset is a whole number
            if (fmod(offset, 1) == 0) {
                
                // Add the offset to the starting time in seconds
                interval.tv_sec = start.tv_sec + ((int)offset);

                // If the offset is not a whole number
            } else {
                
                double secondsOffset;
                // Split the offset into an integral and a fractional part
                double nanosecondsOffset = modf(offset, &secondsOffset) * 1e+9;

                // If the nanoseconds offset would overflow
                if (start.tv_nsec + nanosecondsOffset >= 1e+9) {
                    
                    // Add the offset to the starting time in seconds plus 1
                    interval.tv_sec = start.tv_sec + (int)secondsOffset + 1;
                    // Add the offset (minus 1 second) to the starting time in
                    // nanoseconds
                    interval.tv_nsec =
                        ((start.tv_nsec + (int)nanosecondsOffset) - 1e+9);

                    // Otherwise
                } else {
                    
                    // Add the offset to the starting time in respective
                    // measurements
                    interval.tv_sec = start.tv_sec + (int)secondsOffset;
                    interval.tv_nsec = start.tv_nsec + (int)nanosecondsOffset;
                    
                }
                
            }

            // Pause thread until specified time
            pthread_cond_timedwait(&cond, &mutexLock, &interval);
            
        }

        // Create a buffer of data to send as a ping
        char *pingData = malloc(threadArgs->optValues->s);
        memset(pingData, 1, threadArgs->optValues->s);

        // Update numPings
        numPings++;
        // Allocate space for current ping in ping statistics storage array
        pings[i - 1] = malloc(sizeof(ping_t));

        // Store ping start time
        clock_gettime(CLOCK_REALTIME, &pings[i - 1]->sentTime);

        // Send ping data
        ssize_t numBytes =
            sendto(*threadArgs->socket, pingData, threadArgs->optValues->s, 0,
                   threadArgs->serverAddress->ai_addr,
                   threadArgs->serverAddress->ai_addrlen);
        // Error checking for sendto()
        if (numBytes < 0) {
            
            printErrorAndExit("sendto() failed");

        } else if (numBytes != threadArgs->optValues->s) {
            
            printErrorAndExit(
                "sendto() error: sent unexpected number of bytes");
                
        }
        // Free resources for ping data buffer
        free(pingData);

        // Update the current thread tracker
        currentThread = RECEIVER;

        // Signal remaining threads
        pthread_cond_signal(&cond);
        // Unlock mutex lock
        pthread_mutex_unlock(&mutexLock);
        
    }

    return NULL;
    
}

// Function that performs client ping receiving logic
void *clientReceiver(void *threadArg) {
    
    // Convert and store thread argument as struct object of thread argument
    // values
    threadArgs_t *threadArgs = threadArg;

    // Create an object representing the server address
    struct sockaddr_storage fromAddress;
    // Store the length of the server address information (somehow
    // required for recvfrom() to work)
    socklen_t fromAddressLength = sizeof(fromAddress);

    // Loop through the number of pings to be received
    for (int i = 1; i <= threadArgs->optValues->c; i++) {
        
        // Lock mutex lock for global variable safety
        pthread_mutex_lock(&mutexLock);

        // Loop while the receiver thread is not the active thread
        while (currentThread != RECEIVER) {
            
            // Pause thread until condition and mutex locks are unlocked
            pthread_cond_wait(&cond, &mutexLock);
            
        }

        // Create a buffer for the received ping data
        char buffer[threadArgs->optValues->s];

        // Receive ping response from server
        ssize_t numBytes =
            recvfrom(*threadArgs->socket, buffer, threadArgs->optValues->s, 0,
                     (struct sockaddr *)&fromAddress, &fromAddressLength);
        // Error checking for recvfrom()
        if (numBytes < 0) {
            
            printErrorAndExit("recvfrom() failed");

        } else if (numBytes != threadArgs->optValues->s) {
            
            printErrorAndExit(
                "recvfrom() error: sent unexpected number of bytes");

        } else {
            
            // Mark ping as received
            pings[i - 1]->received = true;
            // Store time ping was received
            clock_gettime(CLOCK_REALTIME, &pings[i - 1]->receivedTime);
            // Calculate and store ping RTT
            pings[i - 1]->difference =
                calculateTimeDifference(&pings[i - 1]->sentTime,
                                        &pings[i - 1]->receivedTime) *
                1e+3;

            // If printing individual ping statistics is allowed
            if (!threadArgs->optValues->n) {
                
                // Print ping statistics
                fprintf(stdout, "\t%d\t%zu\t%.3f\n", i, numBytes,
                        pings[i - 1]->difference);
                        
            }
            
        }

        // Update the current thread tracker
        currentThread = SENDER;

        // Unlock mutex lock
        pthread_mutex_unlock(&mutexLock);
        
    }

    // If we've made it this far, the client is ready to exit
    exitClient();

    return NULL;
    
}

// Function that closes the server socket, prints summary statistics, and exits
// the program with a normal return code
void exitClient() {
    
    // Store current time as ping end time
    struct timespec end;
    clock_gettime(CLOCK_REALTIME, &end);

    // Close open socket
    close(serverSocket);

    // Variable to cound the number of packets that were received
    int numReceived = 0;
    // Variables to track the highest, lowest, and average ping return times
    double min = DBL_MAX;
    double max = 0.0;
    double average;

    // Loop through each ping that was sent
    for (int i = 0; i < numPings; i++) {
        
        // If the ping was marked as received
        if (pings[i]->received) {
            
            // Increment numReceived
            numReceived++;

            // Store the calculated ping return time
            double pingRTT = pings[i]->difference;

            // Add the RTT to the running average
            average += pingRTT;
            // If the RTT is less than the current minimum, update the minimum
            if (pingRTT < min) {
                
                min = pingRTT;
                
            }
            // If the RTT is more than the current maximum, update the maximum
            if (pingRTT > max) {
                
                max = pingRTT;
                
            }
            
        }
        
    }

    // Calculate the average
    average /= numReceived;

    // Print summary statistics
    fprintf(stdout, "\n");
    fprintf(stdout,
            "%d packets transmitted, %d received, %.3f%% packet loss, time "
            "%.0f ms\n",
            numPings, numReceived, (1 - ((float)numReceived / numPings)),
            (calculateTimeDifference(&start, &end) * 1e+3));
    fprintf(stdout, "rtt min/avg/max = %.3f/%.3f/%.3f msec\n", min, average,
            max);

    // Free allocated memory for ping statistic storage
    for (int i = 0; i < numPings; i++) {
        
        free(pings[i]);
        
    }
    free(pings);

    exit(0);
    
}

// Function that performs server ping receiving logic
void *serverReceiver(void *threadArg) {
    
    // Convert and store thread argument as struct object of thread argument
    // values
    threadArgs_t *threadArgs = threadArg;

    // Lock mutex lock for global variable safety
    pthread_mutex_lock(&mutexLock);

    // Loop while the receiver thread is not the active thread
    while (currentThread != RECEIVER) {
        
        // Pause thread until condition and mutex locks are unlocked
        pthread_cond_wait(&cond, &mutexLock);
        
    }

    // Create a buffer to store received ping data
    char buffer[MAX_DATAGRAM_SIZE];

    // Receive ping packet from client
    ssize_t numBytesReceived =
        recvfrom(serverSocket, buffer, MAX_DATAGRAM_SIZE, 0,
                 (struct sockaddr *)threadArgs->clientAddress,
                 threadArgs->clientAddressLength);
    // Error checking for recvfrom()
    if (numBytesReceived < 0) {
        
        printErrorAndExit("recvfrom() failed");
        
    }

    // Update size of client ping for use in server sender thread
    threadArgs->optValues->s = numBytesReceived;

    // Update the current thread tracker
    currentThread = SENDER;

    // Unlock mutex lock
    pthread_mutex_unlock(&mutexLock);
    // Signal remaining threads
    pthread_cond_signal(&cond);

    return NULL;
    
}

// Function that performs server ping sending logic
void *serverSender(void *threadArg) {
    
    // Convert and store thread argument as struct object of thread argument
    // values
    threadArgs_t *threadArgs = threadArg;

    // Lock mutex lock for global variable safety
    pthread_mutex_lock(&mutexLock);

    // Loop while the sender thread is not the active thread
    while (currentThread != SENDER) {
        
        // Pause thread until condition and mutex locks are unlocked
        pthread_cond_wait(&cond, &mutexLock);
        
    }

    // Create a buffer of data to send as a ping
    char *pingData = malloc(threadArgs->optValues->s);
    memset(pingData, 1, threadArgs->optValues->s);

    // Send the received packet back to the client
    ssize_t numBytesSent =
        sendto(serverSocket, pingData, threadArgs->optValues->s, 0,
               (struct sockaddr *)threadArgs->clientAddress,
               *threadArgs->clientAddressLength);
    // Error checking for sendto()
    if (numBytesSent < 0) {
        
        printErrorAndExit("sendto() failed");
        

    } else if (numBytesSent != threadArgs->optValues->s) {
        
        printErrorAndExit("sendto() error: sent unexpected number of bytes");
        
    }
    // Free resources for ping data buffer
    free(pingData);

    // Update the current thread tracker
    currentThread = RECEIVER;

    // Unlock mutex lock
    pthread_mutex_unlock(&mutexLock);
    // Signal remaining threads
    pthread_cond_signal(&cond);

    return NULL;
    
}

// Function that prints statistics about pings, closes the client socket, and
// exits the program with a normal return code
void exitServer() {
    
    // Close open socket
    close(clientSocket);

    exit(0);
    
}