#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <string.h>
#include <stdbool.h>

// Struct for malloc address/size pairs, used to create array
typedef struct mallocPair {

    char *address;
    size_t size;

} mallocPair_t;

// Function prototype for assembleProcessCommand
char *assembleProcessCommand(int argc, char **argv, char *env, char *redir);

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

    // Check that at least one command-line argument is passed
    assert(argc > 1);

    // Assemble large command string to pass to popen
    char env[] = "LD_PRELOAD=./memory_shim.so";
    char redir[] = "2>leaks.db";
    char *command = assembleProcessCommand(argc, argv, env, redir);

    // Open assembled command with popen
    FILE *process = popen(command, "r");
    // Popen error checking
    if (process == NULL) {
        
        perror("Process open failure");
    
    }
    // Close process
    pclose(process);

    // Free memory for command string
    free(command);

    // Open leaks.db file which contains redirected output from malloc/free shim
    FILE *leaksDB = fopen("leaks.db", "r");
    // Check that file is opened correctly
    assert(leaksDB != NULL);

    // Create counter for number of malloc calls
    int numMallocs = 0;
    // Create pointer for array of malloc calls
    mallocPair_t *mallocs = NULL;

    // Reset leaks.db
    rewind(leaksDB);

    // Loop until the end of leaks.db
    while (!feof(leaksDB)) {

        // Create buffer strings for addresses and sizes, initialize with
        // empty string
        char address[1000];
        strncpy(address, "", 1000);
        char size[1000];
        strncpy(size, "", 1000);

        // If fscanf can properly take in delimited fields for address and size
        if (fscanf(leaksDB, "%[^,],%[^\n]\n", address, size) == 2) {

            // Create a malloc pair
            mallocPair_t pair;
            // Allocate memory for address string
            pair.address = malloc(sizeof(char) * strlen(address));
            // Copy local address to malloc pair address
            strcpy(pair.address, address);
            // Convert size to integer/size_t and store in malloc pair size
            pair.size = (size_t)atoi(size);

            // Increment numMallocs
            numMallocs++;
            // If this is the first malloc pair, allocate initial memory for
            // mallocs array
            if (numMallocs == 1) {

                mallocs = malloc(sizeof(mallocPair_t));

            // If this is not the first malloc pair, expand mallocs array
            } else {

                mallocs = realloc(mallocs, (sizeof(mallocPair_t) * numMallocs));

            }
            // Copy the local malloc pair to mallocs array
            mallocs[numMallocs - 1] = pair;

          // Otherwise, read the line from leaks.db normally and print to 
          // stderr (not an address/size pair)
        } else {

            fscanf(leaksDB, "%s", address);
            fprintf(stderr, "%s", address);

        }

    }

    // Close leaks.db
    fclose(leaksDB);

    // Create counter for number of leaks
    int numLeaks = 0;
    // Create size_t array for leaks
    size_t *leaks = NULL;

    // Loop through the number of mallocs
    for (int i = 0; i < numMallocs; i++) {

        // Get the malloc pair address at current array index
        char *tempAddress1 = malloc(sizeof(char) * strlen(mallocs[i].address));
        strcpy(tempAddress1, mallocs[i].address);

        // Assume a match is not found at first
        bool matchFound = false;

        // Loop through number of mallocs from next index
        for (int j = (i + 1); j < numMallocs; j++) {

            // Get the malloc pair address at next array index
            char *tempAddress2 = malloc(sizeof(char) * 
                                        strlen(mallocs[j].address));
            strcpy(tempAddress2, mallocs[j].address);

            // If the two malloc pair addresses match, a match is found so
            // stop looping
            if (strcmp(tempAddress1, tempAddress2) == 0) {

                matchFound = true;
                // Free memory for tempAddress2
                free(tempAddress2);
                break;

            }

        }

        // Free memory for tempAddress1
        free(tempAddress1);

        // If an address match was not found
        if (!matchFound) {

            // If the size of the malloc call is not -1 (used in free() shim, 
            // shouldn't be possible otherwise)
            if (mallocs[i].size != -1) {

                // Increment numLeaks
                numLeaks++;
                // If this is the first leak, allocate initial memory for leaks
                // array
                if (numLeaks == 1) {

                    leaks = malloc(sizeof(size_t));

                  // If this is not the first leak, expand leaks array
                } else {

                    leaks = realloc(leaks, (sizeof(size_t) * numLeaks));

                }
                // Add leak to leaks array
                leaks[numLeaks - 1] = mallocs[i].size;

            }

        }

    }

    // Free memory for mallocs array
    free(mallocs);

    // Create variable to hold total size of leaks
    size_t total = 0;
    // Loop through the number of leaks
    for (int i = 0; i < numLeaks; i++) {

        // Print leak information to stderr
        fprintf(stderr, "LEAK\t%zu\n", leaks[i]);
        // Add leak size to total
        total += leaks[i];

    }
    // Print total leaks information to stderr
    fprintf(stderr, "TOTAL\t%d\t%zu\n", numLeaks, total);

    // Free memory for leaks array
    free(leaks);

    return 0;

}

// Function to concat one large string for the command to be passed to popen
// Takes arguments from parent program run, custom env variable, custom pipe
char *assembleProcessCommand(int argc, char **argv, char *env, char *redir) {
    
    // Assemble the length of all the parts of the command
    int commandLength = strlen(env); // Start with env variable
    commandLength += 1; // Add room for space after env variable
    // Loop through argv excluding original executable call
    for (int i = 1; i < argc; i++) {

        // Add length of each argv plus 1 to make room for space after argv[i]
        commandLength += strlen(argv[i]);
        commandLength += 1;

    }
    // Add length of pipe redirect string
    commandLength += strlen(redir);

    // Allocate memory for command
    char *command = malloc(sizeof(char) * commandLength);
    // Copy empty string into command to start
    strncpy(command, "", commandLength);

    // Concat env plus space after
    strcat(command, env);
    strcat(command, " ");
    // Loop through argv excluding original executable call
    for (int i = 1; i < argc; i++) {

        // Concat argv[i] and a space after
        strcat(command, argv[i]);
        strcat(command, " ");

    }
    // Concat pipe redirect
    strcat(command, redir);

    return command;

}