#include <divine.h>
#include <divine/problem.h>
#include <mpi.h>
#include <cstdlib>

// TODO desc: this is expected to be set in model (__attribute__((constructor)))
// if desired
int _mpi_num_processes = 0;
int _mpi_enable_mpi_comm_self = false;

int _mpi_exit_barrier = 0;

extern "C" {
    int main(...);
    struct global_ctor { int prio; void (*fn)(); };

    int _divine_invoke_main( int mainproto )
    {
        int r;
        switch ( mainproto ) {
            case 0: r = main(); break;
            case 1: r = main( 0 ); break;
            case 2: r = main( 0, 0 ); break;
            case 3: r = main( 0, 0, 0 ); break;
            default: __divine_problem( 1, "don't know how to run main()" );
        }
        return r;
    }

    void _divine_invoke_main_mpi( void *mainproto )
    {
        // run main
        int r = _divine_invoke_main( *((int *) mainproto) );

        // check whether MPI_Finalize has been called
        int finalized; MPI_Finalized( &finalized );
        if ( !finalized ) {
            __divine_problem( Other,
                    "MPI_Finalize not called before program termination" );
        }

        // synchronize thread exits because it's needed apparently TODO desc
        __divine_interrupt_mask();
        ++_mpi_exit_barrier;

        while ( !(_mpi_exit_barrier == _mpi_num_processes &&
                __divine_get_tid() == 0) ) {
            __divine_interrupt_unmask();
            __divine_interrupt_mask();
        }

        exit( r ); // only thread 0 gets here
    }

    void _divine_start( int ctorcount, void *ctors, int mainproto )
    {
        __divine_interrupt_mask(); //FIXME remove
        global_ctor *c = reinterpret_cast< global_ctor * >( ctors );
        for ( int i = 0; i < ctorcount; ++i ) {
            c->fn();
            c ++;
        }
        __divine_interrupt_unmask(); //FIXME remove

        if ( _mpi_num_processes == 0 ) {
            // normal startup (no MPI): run main in one thread
            int r = _divine_invoke_main( mainproto );
            exit( r );
        } else {
            // MPI startup: run main in the specified number of threads
            // (emulating MPI processes)
            __divine_interrupt_mask();

            for ( int i = 0; i < _mpi_num_processes - 1; ++i ) {
                __divine_new_thread( _divine_invoke_main_mpi,
                        (void *) &mainproto );
            }

            // global MPI initialization is run after the process threads have
            // been started, but before they actually had a chance to run
            // (ensured by interrupt masking); this is done in order to preserve
            // sequential numbering of the process threads, since one other
            // thread (the message daemon of MPI_COMM_WORLD) is created as a
            // part of the initialization
            _divine_mpi_global_initialize();
            __divine_interrupt_unmask();

            _divine_invoke_main_mpi( (void *) &mainproto );
        }
    }
}