// The basis of this demo is from Gemini, apparwently from andy-pearce.com
// I have modified it heavly

#include <iostream>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <sys/wait.h>
#include <unistd.h>

using namespace std;

extern "C" {
    void SPINLOCK_WAIT(int * mutex);
    void SPINLOCK_WAIT2(int * mutex);
    void SPINLOCK_UNLOCK(int * mutex);
}


struct SharedDataT {
    int counter;
    int mutex;
};

void DoWork(SharedDataT * shared, int iterations, bool doMutex) ;
void Usage(string name);

int main(int argc, char * argv[]) {
    bool doMutex {true};
    int processes{5};
    int iterations{1000};
    bool argsError{false};
    bool verbose{false};

    int tmp;
    int opt;

    while ((opt = getopt(argc, argv, "vhmp:i:")) != -1) {
        switch(opt) {
            case 'v':
                verbose = true;
                break;
            case 'h': 
                argsError = true;
                break;
            case 'm':
                doMutex = false;
                break;
            case 'p':
                tmp = stoi(optarg);
                if (tmp > 1) {
                    processes = tmp;
                }
                break;
            case 'i':
                tmp = stoi(optarg);
                if (tmp > 1) {
                    iterations = tmp;
                }
                break;
            default:
                argsError = true;
        }
    }

    if (argsError) {
        Usage(argv[0]);
    }

    if (verbose) {
        cout << endl;
        cout << "Do Mutex " << boolalpha << doMutex << endl;
        cout << "Processes: " << processes << endl;
        cout << "Iterations: " << iterations << endl;
        cout << endl;
    }

    int prot {PROT_READ | PROT_WRITE};
    int flags{MAP_SHARED| MAP_ANONYMOUS};

    SharedDataT * shared = static_cast<SharedDataT *>(mmap(NULL, 
                         sizeof(struct SharedDataT), prot, flags, -1, 0));

    if (shared == nullptr) {
        perror("mmap: ");
        return 0;
    }
    shared->counter = 0;
    shared->mutex = 0;

    for (int i = 0; i < processes; ++i) {
        if (fork() == 0) {
            DoWork(shared, iterations, doMutex);
            return 0;
        }
    }

    // wait for all workers to finsh 
    for (int i = 0; i < processes; ++i){
       wait(NULL);
    }

    cout << "The count is " << shared->counter << endl;
    cout << "It should be " << processes * iterations << endl;

    munmap(shared, sizeof(SharedDataT));
    return 0;
}

void Usage(string name) {
    cout << endl;
    cout << "Usage " << name << endl;
    cout <<"\t-m   ....... do not use mutex" << endl;
    cout <<"\t-i n ....... exch process should perform n iterations" << endl;
    cout <<"\t-p n ....... start p processes" << endl;
    cout <<"\t-v   ....... print configuration status" << endl;
    cout <<"\t-h   ....... print help" << endl;
    cout << endl;
}

void DoWork(SharedDataT * shared, int iterations, bool doMutex) {
   for (int j = 0; j < iterations; ++j) {
         // usleep will force a contex switch
         usleep(rand() % 10);

         // if we need to, acquire the lock
         if(doMutex) {
             SPINLOCK_WAIT2(&(shared->mutex));
         }

         shared->counter++;

         // if we need to, release the lock
         if (doMutex) {
             SPINLOCK_UNLOCK(&(shared->mutex));
         }
     }
}