#include <stdio.h>
#include <assert.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>
#include <unistd.h>
#include <sys/ptrace.h>
#include <sys/wait.h>
#include <sys/reg.h>
#include <sys/types.h>

// Struct for unique syscalls with syscall number and counter for number 
// of calls to a unique syscall
typedef struct syscall {

    int syscallNum;
    int numCalls;

} syscall_t;

// Function prototypes
char **splitArguments(char *arg);
bool syscallStatus(pid_t child, int status);

int main(int argc, char **argv) {

    // Split second command-line argument
    char **args = splitArguments(argv[1]);

    // Fork process
    pid_t child = fork();
    // If process is child process
    if (child == 0) {

        // Allow child process to be traced
        ptrace(PTRACE_TRACEME);

        // Stop child process to allow parent process to trace
        kill(getpid(), SIGSTOP);

        // Call execvp passing the first of the split arguments and 
        // the whole split arguments array
        execvp(args[0], args);

      // If process is parent process
    } else {
        
        // Make sure there are two arguments passed (one for program to run, 
        // one for output file name)
        assert(argc == 3);
        
        int status;
        int syscallNum;

        // Create counter for number of unique syscalls
        int numUniqueSyscalls = 0;
        // Create a pointer array to store unique syscalls
        syscall_t *syscalls;

        // Wait for the child process to stop itself
        waitpid(child, &status, 0);

        // Use PTRACE_O_TRACESYSGOOD to trace only normal syscalls, not traps
        ptrace(PTRACE_SETOPTIONS, child, 0, PTRACE_O_TRACESYSGOOD);

        // Loop until the loop is broken based on the return value of 
        // syscallStatus()
        // Reference: https://blog.nelhage.com/2010/08/write-yourself-an-strace-in-70-lines-of-code/
        //            (demonstrated that a loop should be used here)
        while (true) {

            // Check if syscallStatus is false (indicating full
            // child program exit)
            if (syscallStatus(child, status) == false) {

                break;

            }
            
            // Store the current syscall number
            syscallNum = ptrace(PTRACE_PEEKUSER, child, (sizeof(long) * ORIG_RAX), NULL);

            bool matchFound = false;
            int syscallIndex;
            
            // Loop through syscalls array
            for (int i = 0; i < numUniqueSyscalls; i++) {

                // If the current syscall is not unique, a match has been 
                // found and store the index where the match was
                if (syscalls[i].syscallNum == syscallNum) {

                    matchFound = true;
                    syscallIndex = i;
                    break;

                }

            }

            // If a match has been found and the syscall is not unique, 
            // increment the number of calls for that syscall
            if (matchFound) {

                syscalls[syscallIndex].numCalls++;

              // If a match has not been found and the syscall is unique
            } else {

                // Increment numUniqueSyscalls
                numUniqueSyscalls++;

                // If this is the first unique syscall, allocate initial 
                // memory for syscalls array
                if (numUniqueSyscalls == 1) {

                    syscalls = malloc(sizeof(syscall_t));

                  // If this is not the first unique syscall, expand syscalls 
                  // array
                } else {

                    syscalls = realloc(syscalls, (sizeof(syscall_t) * numUniqueSyscalls));

                }

                // Set the new unique syscall number and increment the number 
                // of calls to 1
                syscalls[numUniqueSyscalls - 1].syscallNum = syscallNum;
                syscalls[numUniqueSyscalls - 1].numCalls++;

            }

            // Check syscallStatus for full program exit again (occurs 
            // after exit_group() system call)
            if (syscallStatus(child, status) == false) {

                break;

            }

        }
        
        // Open output file whose name is specified as a command-line argument
        FILE *output = fopen(argv[2], "w+");
        // File error checking
        assert(output != NULL);
        
        // Loop through syscalls array
        for (int i = 0; i < numUniqueSyscalls; i++) {

            // Print output to output file
            fprintf(output, "%d\t%d\n", syscalls[i].syscallNum, syscalls[i].numCalls);

        }

        // Free args array
        free(args);
        // Free syscalls array
        free(syscalls);

        // Close output file
        fclose(output);

    }

    return 0;

}

// Function to split arguments in argv[1] so that they can be properly 
// passed to execvp()
char **splitArguments(char *arg) {

    int numSpaces = 0;
    // Loop through arg passed to function
    for (int i = 0; i < strlen(arg); i++) {

        // If character is a space, increment numSpaces
        if (arg[i] == ' ') {

            numSpaces++;

        }

    }

    // If the first character in the arg passed to a function is a 
    // space, re-build the array without the leading space and decrement 
    // numSpaces so that we properly null-terminate the array
    if (arg[0] == ' ') {

        // Create a temp string
        char *temp = malloc(sizeof(char) * (strlen(arg) - 1));

        // Copy everything after the first character
        for (int i = 1; i < strlen(arg); i++) {

            temp[i - 1] = arg[i];

        }

        // Assign the temp string to the original
        arg = temp;

        numSpaces--;

    }

    // If the last character in the arg passed to function is a 
    // space, decrement numSpaces so that we properly null-terminate the array
    if (arg[strlen(arg) - 1] == ' ') {

        numSpaces--;

    }

    // If there are no spaces, we already have our full arg string
    if (numSpaces == 0) {

        // Create and allocate a double pointer to store arg and NULL
        char **args = malloc(sizeof(char *) * 2);
        
        // Allocate first index of double pointer
        args[0] = malloc(sizeof(char) * strlen(arg));
        // Add each character from arg to first index of args
        for (int i = 0; i < strlen(arg); i++) {

            args[0][i] = arg[i];

        }

        // Set second index of args to NULL
        args[1] = NULL;

        return args;

    } else {

        // Create and allocate a double pointer to store space-separated 
        // args and NULL at the end
        char **args = malloc(sizeof(char *) * (numSpaces + 2));

        int numCharacters = 0;
        int stopPos = 0;
        // Loop through number of spaces (plus 1 since there is an argument 
        // before the first space)
        for (int i = 0; i < (numSpaces + 1); i++) {

            // Loop through the characters in arg
            for (int j = stopPos; j < strlen(arg); j++) {

                // If character is not a space, increment numCharacters
                if (arg[j] != ' ') {

                    numCharacters++;

                  // If character is a space, set the new stop position and break 
                  // out of inner loop
                } else {

                    stopPos = j;
                    break;

                }

            }

            // Allocate the current index of args based on the number of 
            // characters, then reset numCharacters
            args[i] = malloc(sizeof(char) * numCharacters);
            numCharacters = 0;

        }

        // Loop through the characters in arg
        for (int i = 0, j = 0, k = 0; i < strlen(arg); i++) {

            // If character is not a space
            if (arg[i] != ' ') {

                // Add character to current index in args
                args[j][k] = arg[i];
                k++;

              // If character is a space
            } else {

                // Move to the next index in args
                j++;
                k = 0;

            }

        }

        // Set last index of args to NULL
        args[numSpaces + 2] = NULL;

        return args;

    }

}

// Function to check if the child process has stopped (either temporarily 
// or permanently)
// Reference: https://blog.nelhage.com/2010/08/write-yourself-an-strace-in-70-lines-of-code/
//            (demonstrated that a loop should be used here)
bool syscallStatus(pid_t child, int status) {

    // Loop infinitely, loop will be stopped by return statement
    while (true) {

        // Trace child syscall
        ptrace(PTRACE_SYSCALL, child, 0, 0);

        // Wait for child process to stop itself
        waitpid(child, &status, 0);

        // If the child has reached SIGSTOP (temporary), return true
        if (WIFSTOPPED(status) && WSTOPSIG(status) & 0x80) {

            return true;

        }

        // If the child process has stopped permanently, return false
        if (WIFEXITED(status)) {

            return false;

        }

    }

}