#include <mpi.h>
#include <iostream>
#include <thread>
#include <chrono>
#include <vector>

namespace logger {
std::ostream& error() {
    return std::cerr;
}
std::ostream& verbose() {
    return std::cout;
}
}

void check(int const mpiErrorCode)
{
    if (mpiErrorCode != MPI_SUCCESS) {
        logger::error() << "MPI Operation failed with code: " << mpiErrorCode << "\n";
        std::exit(2);
    }
}

void runSendBig2Barrier(MPI_Comm comm, int count)
{
    int rank, size;
    check(MPI_Comm_rank(comm, &rank));
    check(MPI_Comm_size(comm, &size));
    
    int constexpr TAG = 1234;
    std::vector<int> payload(count);
    MPI_Request req;

    if (rank == 1) {
        MPI_Isend(payload.data(), count, MPI_INT, 0, TAG, comm, &req);
        logger::verbose() << "Done MPI_Isend\n";
    } else {
        std::this_thread::sleep_for(std::chrono::seconds(1));
        logger::verbose() << "Done sleep_for\n";
    }
    logger::verbose() << "Barrier 1\n";
    check(MPI_Barrier(comm));
    logger::verbose() << "Barrier 2\n";
    check(MPI_Barrier(comm));
    logger::verbose() << "Barriers all done.\n";

    if (rank == 0) {
        check(MPI_Recv(payload.data(), count, MPI_INT, 1, TAG, comm, MPI_STATUS_IGNORE));
        logger::verbose() << "Done MPI_Recv\n";
    } else if (rank == 1) {
        check(MPI_Wait(&req, MPI_STATUS_IGNORE));
    }

    logger::verbose() << "All done in runSendBig2Barrier, about to MPI_Comm_free\n";
}

int main(int argc, char ** argv) {
    check(MPI_Init(&argc, &argv));

    int const count = (argc > 1) ? std::atoi(argv[1]) : 10000000;

    runSendBig2Barrier(MPI_COMM_WORLD, count);
    check(MPI_Finalize());
}