!> @file mpi_runtime.F90
!> @brief Process-global MPI lifecycle and identity queries.
!!
!! Owns the single call to MPI_Init / MPI_Finalize for any executable that
!! includes this module, and exposes the world communicator together with the
!! local rank and total rank count. It runs at any `mpirun -np N` (and at
!! `-np 1`, where the rank/size queries degenerate cleanly).
!!
!! This module uses only the lowest-common-denominator MPI features (Init,
!! Finalize, Comm_rank, Comm_size, Allreduce, COMM_WORLD). The Fortran 2008
!! `mpi_f08` module gives us type-safe handles on every implementation that
!! provides it; MS-MPI on Windows does not ship `mpi_f08` yet, so we fall
!! back to the legacy `mpi` module when `CFD_SOLVER_LEGACY_MPI` is defined
!! at preprocess time (set via `-DCFD_SOLVER_LEGACY_MPI` on the Windows CI
!! job). The call-site code is identical either way because `mpi_world()`
!! and the named constants resolve to the handle type of whichever module
!! is in scope.

module mpi_runtime
#ifdef CFD_SOLVER_LEGACY_MPI
  use mpi, only: MPI_Init, MPI_Finalize, MPI_Initialized, MPI_Abort, &
                 MPI_Comm_rank, MPI_Comm_size, MPI_COMM_WORLD, MPI_PROC_NULL
#else
  use mpi_f08, only: MPI_Init, MPI_Finalize, MPI_Initialized, MPI_Abort, &
                     MPI_Comm_rank, MPI_Comm_size, MPI_COMM_WORLD, &
                     MPI_Comm, MPI_PROC_NULL
#endif
  use, intrinsic :: iso_fortran_env, only: error_unit
  implicit none
  private

  public :: mpi_runtime_init
  public :: mpi_runtime_finalize
  public :: mpi_runtime_is_initialised
  public :: my_rank
  public :: n_ranks
  public :: mpi_world
  public :: parallel_fatal
  public :: mpi_proc_null_value

  !> Cached identity values, populated by mpi_runtime_init.
  integer, save :: g_my_rank = -1
  integer, save :: g_n_ranks = -1
  logical, save :: g_initialised = .false.

contains

  !> Initialise the MPI library and cache rank/size.
  !!
  !! Calling this more than once is a programmer error; the second call aborts
  !! the program via `error stop`. Callers that need to check status should use
  !! `mpi_runtime_is_initialised` first.
  subroutine mpi_runtime_init()
    integer :: ierr
    logical :: already

    if (g_initialised) &
      error stop 'mpi_runtime: mpi_runtime_init called more than once'

    call MPI_Initialized(already, ierr)
    if (ierr /= 0) error stop 'mpi_runtime: MPI_Initialized failed'
    if (already) error stop 'mpi_runtime: MPI was already initialised outside mpi_runtime'

    call MPI_Init(ierr)
    if (ierr /= 0) error stop 'mpi_runtime: MPI_Init failed'

    call MPI_Comm_rank(MPI_COMM_WORLD, g_my_rank, ierr)
    if (ierr /= 0) error stop 'mpi_runtime: MPI_Comm_rank failed'
    call MPI_Comm_size(MPI_COMM_WORLD, g_n_ranks, ierr)
    if (ierr /= 0) error stop 'mpi_runtime: MPI_Comm_size failed'

    g_initialised = .true.
  end subroutine mpi_runtime_init

  !> Finalise the MPI library. Idempotent: a second call is a no-op.
  subroutine mpi_runtime_finalize()
    integer :: ierr
    logical :: was_init

    if (.not. g_initialised) return
    call MPI_Initialized(was_init, ierr)
    if (was_init) then
      call MPI_Finalize(ierr)
      if (ierr /= 0) write (error_unit, '(A)') 'mpi_runtime: MPI_Finalize returned non-zero'
    end if
    g_initialised = .false.
    g_my_rank = -1
    g_n_ranks = -1
  end subroutine mpi_runtime_finalize

  !> True after `mpi_runtime_init` and before `mpi_runtime_finalize`.
  pure function mpi_runtime_is_initialised() result(yes)
    logical :: yes
    yes = g_initialised
  end function mpi_runtime_is_initialised

  !> Local rank in MPI_COMM_WORLD. Aborts if called before init.
  function my_rank() result(r)
    integer :: r
    if (.not. g_initialised) error stop 'mpi_runtime: my_rank called before init'
    r = g_my_rank
  end function my_rank

  !> Size of MPI_COMM_WORLD. Aborts if called before init.
  function n_ranks() result(n)
    integer :: n
    if (.not. g_initialised) error stop 'mpi_runtime: n_ranks called before init'
    n = g_n_ranks
  end function n_ranks

  !> Return the world communicator handle. Convenience accessor so call sites
  !! never reference MPI_COMM_WORLD directly. Always call this function at the
  !! point of use rather than caching the return value, so the call site stays
  !! correct once sub-communicators are introduced.
  function mpi_world() result(comm)
#ifdef CFD_SOLVER_LEGACY_MPI
    integer :: comm        !< Legacy `mpi` module: handles are bare integers.
#else
    type(MPI_Comm) :: comm !< `mpi_f08` module: derived-type handle.
#endif
    comm = MPI_COMM_WORLD
  end function mpi_world

  !> Expose MPI_PROC_NULL as a value so callers that don't `use mpi*`
  !! themselves can compare against it.
  pure function mpi_proc_null_value() result(v)
    integer :: v
    v = MPI_PROC_NULL
  end function mpi_proc_null_value

  !> Abort the entire MPI world with a diagnostic message.
  !!
  !! Use this for fatal errors that may be raised on a subset of ranks. The
  !! message is printed by the calling rank; `MPI_Abort` then terminates every
  !! rank. At `-np 1` this is functionally equivalent to `error stop`, while
  !! remaining correct when a fatal error is raised on only a subset of ranks.
  subroutine parallel_fatal(message)
    character(len=*), intent(in) :: message
    integer :: ierr

    write (error_unit, '(A,I0,A,A)') '[rank=', g_my_rank, '] FATAL: ', trim(message)
    flush (error_unit)
    call MPI_Abort(MPI_COMM_WORLD, 1, ierr)
    ! Should not return; if it does, fall through to error stop.
    error stop 'mpi_runtime: MPI_Abort returned without terminating'
  end subroutine parallel_fatal

end module mpi_runtime
