// Copyright (C) 2010, Guy Barrand. All rights reserved.
// See the file tools.license for terms.

#ifndef toolx_mpi_wait_buffer
#define toolx_mpi_wait_buffer

#include <mpi.h>

#include <ostream>

namespace toolx {
namespace mpi {

inline bool wait_buffer(std::ostream& a_out,int a_rank,int a_src,int a_tag,const MPI_Comm& a_comm,
                        int& a_buffer_size,char*& a_buffer,int& a_probe_src,bool a_verbose = false) {
  a_buffer = 0;
  a_buffer_size = 0;
  a_probe_src = -1;

  MPI_Status status;
  if(::MPI_Probe(a_src,a_tag,a_comm,&status)!=MPI_SUCCESS) {
    a_out << "toolx::mpi::wait_buffer : rank " << a_rank << " : MPI_Probe : failed." << std::endl;
    return false;
  }

  if(::MPI_Get_count(&status,MPI_CHAR,&a_buffer_size)!=MPI_SUCCESS) {
    a_out << "toolx::mpi::wait_buffer : rank " << a_rank << " : MPI_Get_count : failed." << std::endl;
    a_buffer_size = 0;
    return false;
  }

  a_probe_src = status.MPI_SOURCE;

  if(!a_buffer_size) {
    a_out << "exlb::mpi::wait_buffer : MPI_Get_count returns zero data." << std::endl;
    a_probe_src = -1;
    return false;
  }

  if(a_verbose) a_out << "toolx::mpi::wait_buffer : rank " << a_rank << " : get_count " << a_buffer_size << std::endl;

  a_buffer = new char[a_buffer_size];
  if(!a_buffer) {
    a_out << "toolx::mpi::wait_buffer : rank " << a_rank << " : can't alloc buffer of size " << a_buffer_size << std::endl;
    a_buffer_size = 0;
    a_probe_src = -1;
    return false;
  }

  if(::MPI_Recv(a_buffer,a_buffer_size,MPI_CHAR,status.MPI_SOURCE,status.MPI_TAG,a_comm,&status)!=MPI_SUCCESS) {
    a_out << "toolx::mpi::wait_buffer : rank " << a_rank << " : MPI_Recv : failed." << std::endl;
    a_buffer_size = 0;
    delete [] a_buffer;
    a_buffer = 0;
    a_probe_src = -1;
    return false;
  }

  if(a_verbose) a_out << "toolx::mpi::wait_buffer : rank " << a_rank << " : unpack data ..." << std::endl;

  return true;
}

}}

#endif