#include <upcxx/upcxx.hpp>
#include <thread>
#include <stdio.h>
#include <assert.h>

using namespace std;

struct Data {
    int data;
    bool done;
    Data(int d, bool b) : data(d), done(b) {};
};
using dobj_d_t = upcxx::dist_object<Data*>;

int main(int argc, char** argv)
{
    upcxx::init();
    int N = atoi(argv[1]);
    cout << "Running " << N << " times" << endl;
    Data* data = new Data(-1, false);
    dobj_d_t local_d(data);
    const int n_ranks = upcxx::rank_n();
    const int rank    = upcxx::rank_me();
    int dest          = (rank + 1) % n_ranks;
    int expected      = rank;

    upcxx::liberate_master_persona();

    for(int n = 0; n < N; n++) {

        (*data) = Data(-1, false);

        thread tc = thread([&](){
            upcxx::persona_scope scope(upcxx::master_persona());
            printf("Comm %d started -> %d\n", rank, dest);
            upcxx::rpc_ff(dest,
                [](dobj_d_t &data, int payload) {
                    printf("Rank %d received RPC, data is at %p\n", upcxx::rank_me(), *data);
                    (*data)->data = payload;
                    (*data)->done = true;
                    printf("Rank %d received data %d\n", upcxx::rank_me(), payload);
                },
                local_d, dest);
            printf("Comm %d done\n", rank);
            while(! data->done) {
                upcxx::progress();
            }
            printf("Progress %d done\n", rank);
        });

        tc.join();

        if(data->data == expected) {
            printf("SUCCESS\n");
        } else {
            printf("FAILED\n");
        }

    }

    upcxx::persona_scope scope(upcxx::master_persona());
    upcxx::barrier();
    upcxx::finalize();
    delete data;
    return 0;
}