#pragma once

#include <cassert>
#include <mpi.h>

class MpiEnvironment { // wrapper for creating and destroying the environment
	std::size_t worldRank_{std::numeric_limits<std::size_t>::max()};
	std::size_t worldSize_{0};
	bool isMaster_{false};

  public:
	std::size_t worldRank() const { return worldRank_; }
	std::size_t worldSize() const { return worldSize_; }
	bool isMaster() const { return isMaster_; }

	friend void swap(MpiEnvironment& first, MpiEnvironment& second) noexcept {
		using std::swap;
		swap(first.worldRank_, second.worldRank_);
		swap(first.worldSize_, second.worldSize_);
		swap(first.isMaster_, second.isMaster_);
	}

	MpiEnvironment(int& argc, char* argv[], int required_thread_support) {
		int provided_thread_support;
		MPI_Init_thread(&argc, &argv, required_thread_support, &provided_thread_support);
		assert(required_thread_support == provided_thread_support);

		worldRank_ = [] { int r;  MPI_Comm_rank(MPI_COMM_WORLD, &r); return static_cast<std::size_t>(r); }();
		worldSize_ = [] {int s;  MPI_Comm_size(MPI_COMM_WORLD, &s); return static_cast<std::size_t>(s); }();
		isMaster_ = {worldRank_ == 0};
	}
	~MpiEnvironment() {
		if (worldRank_ != std::numeric_limits<std::size_t>::max()) { MPI_Finalize(); }
	}
	MpiEnvironment(MpiEnvironment&) = delete;
	MpiEnvironment& operator=(MpiEnvironment&) = delete;
	MpiEnvironment(MpiEnvironment&& other) noexcept { swap(*this, other); }
	MpiEnvironment& operator=(MpiEnvironment&& other) noexcept {
		swap(*this, other);
		return *this;
	}
};
