!--------------------------------------------------------------------------------------------------!
!   CP2K-RS MPI Interface: Wrapper for MPI-aware initialization                                   !
!   Copyright 2025 CP2K-RS developers                                                             !
!                                                                                                  !
!   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
!--------------------------------------------------------------------------------------------------!

! **************************************************************************************************
!> \brief MPI-aware initialization wrapper for CP2K when MPI is already initialized by host
!> \par History
!>       12.2024 created as part of cp2k-rs MPI integration
!> \author CP2K-RS Developers
!> \note This module provides a wrapper to initialize CP2K when MPI has already been initialized
!>       by the host application (e.g., mpi4py in Python, mpi crate in Rust). It checks if MPI
!>       is already initialized and calls the appropriate CP2K initialization function.
! **************************************************************************************************
MODULE libcp2k_mpi
   USE ISO_C_BINDING,                   ONLY: C_INT
   USE f77_interface,                   ONLY: init_cp2k,&
                                              finalize_cp2k,&
                                              default_para_env
   USE base_hooks,                      ONLY: cp_abort_hook

   IMPLICIT NONE

   PRIVATE

   ! Public interface
   PUBLIC :: cp2k_init_with_mpi_check
   PUBLIC :: cp2k_finalize_with_mpi_check

CONTAINS

! **************************************************************************************************
!> \brief Initialize CP2K with automatic MPI detection
!> \par This function checks if MPI is already initialized (e.g., by mpi4py or Rust's mpi crate)
!>      and calls the appropriate CP2K initialization function. If MPI is already initialized,
!>      it uses the existing MPI environment. Otherwise, it initializes MPI as part of CP2K init.
!> \return success Returns 0 on success, non-zero on failure
! **************************************************************************************************
   FUNCTION cp2k_init_with_mpi_check() RESULT(success) BIND(C, name="cp2k_init_with_mpi_check")
      INTEGER(C_INT)                                     :: success

      INTEGER                                            :: ierr
      LOGICAL                                            :: mpi_is_initialized

#if defined(__parallel)
      ! Check if MPI is already initialized
      CALL MPI_Initialized(mpi_is_initialized, ierr)
      IF (ierr /= 0) THEN
         success = ierr
         RETURN
      END IF

      IF (mpi_is_initialized) THEN
         ! MPI already initialized by host application (mpi4py, mpi crate, etc.)
         ! Use init_cp2k with init_mpi=.FALSE. to reuse existing MPI environment.
         !
         ! NOTE: init_cp2k(.FALSE.) executes `default_para_env = mp_comm_world` which
         ! triggers the mp_comm_assign operator. That operator copies the handle and
         ! calls init(.FALSE.) internally, which:
         !   - queries MPI_COMM_WORLD for num_pe and mepos
         !   - sets ref_count = 1
         !   - sets owns_group = .FALSE. (will NOT call MPI_Comm_free on cleanup)
         ! cp_logger_create then retains default_para_env → ref_count = 2.
         ! We must NOT call init() again here — doing so resets ref_count to 1,
         ! which causes a ref_count underflow abort during finalize_cp2k.
         CALL init_cp2k(.FALSE., ierr)
         IF (ierr /= 0) THEN
            success = ierr
            RETURN
         END IF
         success = 0
      ELSE
         ! MPI not initialized yet, let CP2K initialize it
         CALL init_cp2k(.TRUE., ierr)
         IF (ierr /= 0) THEN
            success = ierr
            RETURN
         END IF
         success = 0
      END IF
#else
      ! Non-parallel build, just initialize without MPI
      CALL init_cp2k(.FALSE., ierr)
      IF (ierr /= 0) THEN
         success = ierr
         RETURN
      END IF
      success = 0
#endif

   END FUNCTION cp2k_init_with_mpi_check

! **************************************************************************************************
!> \brief Finalize CP2K with automatic MPI detection
!> \par This function checks if MPI should be finalized. If CP2K initialized MPI, it will
!>      finalize it. If MPI was already initialized by the host, CP2K won't finalize it
!>      (the host is responsible for that).
!> \return success Returns 0 on success, non-zero on failure
!> \note In practice, we track whether we initialized MPI by checking if it was already
!>       initialized when we started. If it was already initialized, we don't finalize it.
! **************************************************************************************************
   FUNCTION cp2k_finalize_with_mpi_check() RESULT(success) BIND(C, name="cp2k_finalize_with_mpi_check")
      INTEGER(C_INT)                                     :: success

      INTEGER                                            :: ierr
      LOGICAL                                            :: mpi_is_finalized

#if defined(__parallel)
      ! Check if MPI is already finalized
      CALL MPI_Finalized(mpi_is_finalized, ierr)
      IF (ierr /= 0) THEN
         success = ierr
         RETURN
      END IF

      ! Null the abort hook before tearing down the logger/para_env.
      ! During finalize_cp2k, the default logger and para_env are destroyed in
      ! sequence.  If any step triggers cp_abort, the handler tries to log via
      ! the logger that is in the middle of being freed, causing a recursive
      ! SIGSEGV.  By clearing cp_abort_hook here, failures in the cleanup path
      ! fall through to the default Fortran STOP instead of recursing.
      NULLIFY(cp_abort_hook)

      IF (mpi_is_finalized) THEN
         ! MPI already finalized, just finalize CP2K internals without touching MPI
         CALL finalize_cp2k(.FALSE., ierr)
         success = 0
      ELSE
         ! MPI still active, but we should NOT finalize it if we didn't initialize it
         ! The host application (mpi4py/mpi crate) will finalize MPI
         ! So always use finalize_mpi=.FALSE.
         CALL finalize_cp2k(.FALSE., ierr)
         success = 0
      END IF
#else
      NULLIFY(cp_abort_hook)
      ! Non-parallel build, just finalize without MPI
      CALL finalize_cp2k(.FALSE., ierr)
      success = 0
#endif

   END FUNCTION cp2k_finalize_with_mpi_check

END MODULE libcp2k_mpi
