#include <petsc/private/petscimpl.h>

#if defined(PETSC_HAVE_CUDA)
#include <cuda_runtime.h>
#endif

#if defined(PETSC_HAVE_HIP)
#include <hip/hip_runtime.h>
#endif

static PetscInt petsc_checkpointer_intensity = 1;

/*@
   PetscCheckPointerSetIntensity - An intense pointer check registers a signal handler and attempts to dereference to
   confirm whether the address is valid.  An intensity of 0 never uses signal handlers, 1 uses them when not in a "hot"
   function, and intensity of 2 always uses a signal handler.

   Not Collective

   Input Parameter:
.  intensity - how much to check pointers for validity

   Options Database:
.  -check_pointer_intensity - intensity (0, 1, or 2)

   Level: advanced

.seealso: `PetscCheckPointer()`, `PetscFunctionBeginHot()`
@*/
PetscErrorCode PetscCheckPointerSetIntensity(PetscInt intensity) {
  PetscFunctionBegin;
  switch (intensity) {
  case 0:
  case 1:
  case 2: petsc_checkpointer_intensity = intensity; break;
  default: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Intensity %" PetscInt_FMT " not in 0,1,2", intensity);
  }
  PetscFunctionReturn(0);
}

/* ---------------------------------------------------------------------------------------*/

#if defined(PETSC_HAVE_SETJMP_H)
#include <setjmp.h>
static jmp_buf   PetscSegvJumpBuf;
static PetscBool PetscSegvJumpBuf_set;

/*@C
   PetscSignalSegvCheckPointerOrMpi - To be called from a signal handler for SIGSEGV.  If the signal was received
   while executing PetscCheckPointer()/PetscCheckMpiXxxAwareness(), this function longjmps back there, otherwise returns
   with no effect. This function is called automatically by PetscSignalHandlerDefault().

   Not Collective

   Level: developer

.seealso: `PetscPushSignalHandler()`
@*/
void PetscSignalSegvCheckPointerOrMpi(void) {
  if (PetscSegvJumpBuf_set) longjmp(PetscSegvJumpBuf, 1);
}

/*@C
     PetscCheckPointer - Returns PETSC_TRUE if a pointer points to accessible data

   Not Collective

   Input Parameters:
+     ptr - the pointer
-     dtype - the type of data the pointer is suppose to point to

   Level: developer

.seealso: `PetscCheckPointerSetIntensity()`
@*/
PetscBool PetscCheckPointer(const void *ptr, PetscDataType dtype) {
  if (PETSC_RUNNING_ON_VALGRIND) return PETSC_TRUE;
  if (!ptr) return PETSC_FALSE;
  if (petsc_checkpointer_intensity < 1) return PETSC_TRUE;

#if PetscDefined(USE_DEBUG)
  /* Skip the verbose check if we are inside a hot function. */
  if (petscstack.hotdepth > 0 && petsc_checkpointer_intensity < 2) return PETSC_TRUE;
#endif

  PetscSegvJumpBuf_set = PETSC_TRUE;

  if (setjmp(PetscSegvJumpBuf)) {
    /* A segv was triggered in the code below hence we return with an error code */
    PetscSegvJumpBuf_set = PETSC_FALSE;
    return PETSC_FALSE;
  } else {
    switch (dtype) {
    case PETSC_INT: {
      PETSC_UNUSED PetscInt x = (PetscInt) * (volatile PetscInt *)ptr;
      break;
    }
#if defined(PETSC_USE_COMPLEX)
    case PETSC_SCALAR: { /* C++ is seriously dysfunctional with volatile std::complex. */
#if defined(PETSC_USE_CXXCOMPLEX)
      PetscReal                         xreal = ((volatile PetscReal *)ptr)[0], ximag = ((volatile PetscReal *)ptr)[1];
      PETSC_UNUSED volatile PetscScalar x = xreal + PETSC_i * ximag;
#else
      PETSC_UNUSED PetscScalar x = *(volatile PetscScalar *)ptr;
#endif
      break;
    }
#endif
    case PETSC_REAL: {
      PETSC_UNUSED PetscReal x = *(volatile PetscReal *)ptr;
      break;
    }
    case PETSC_BOOL: {
      PETSC_UNUSED PetscBool x = *(volatile PetscBool *)ptr;
      break;
    }
    case PETSC_ENUM: {
      PETSC_UNUSED PetscEnum x = *(volatile PetscEnum *)ptr;
      break;
    }
    case PETSC_CHAR: {
      PETSC_UNUSED char x = *(volatile char *)ptr;
      break;
    }
    case PETSC_OBJECT: {
      PETSC_UNUSED volatile PetscClassId classid = ((PetscObject)ptr)->classid;
      break;
    }
    default:;
    }
  }
  PetscSegvJumpBuf_set = PETSC_FALSE;
  return PETSC_TRUE;
}

#define PetscMPICUPMAwarnessCheckFunction \
  PetscBool PetscMPICUPMAwarenessCheck(void) { \
    cupmError_t cerr = cupmSuccess; \
    int         ierr, hbuf[2] = {1, 0}, *dbuf = NULL; \
    PetscBool   awareness = PETSC_FALSE; \
    cerr                  = cupmMalloc((void **)&dbuf, sizeof(int) * 2); \
    if (cerr != cupmSuccess) return PETSC_FALSE; \
    cerr = cupmMemcpy(dbuf, hbuf, sizeof(int) * 2, cupmMemcpyHostToDevice); \
    if (cerr != cupmSuccess) return PETSC_FALSE; \
    PetscSegvJumpBuf_set = PETSC_TRUE; \
    if (setjmp(PetscSegvJumpBuf)) { \
      /* If a segv was triggered in the MPI_Allreduce below, it is very likely due to the MPI is not GPU-aware */ \
      awareness = PETSC_FALSE; \
    } else { \
      ierr = MPI_Allreduce(dbuf, dbuf + 1, 1, MPI_INT, MPI_SUM, PETSC_COMM_SELF); \
      if (!ierr) awareness = PETSC_TRUE; \
    } \
    PetscSegvJumpBuf_set = PETSC_FALSE; \
    cerr                 = cupmFree(dbuf); \
    if (cerr != cupmSuccess) return PETSC_FALSE; \
    return awareness; \
  }

#if defined(PETSC_HAVE_CUDA)
#define cupmError_t                cudaError_t
#define cupmMalloc                 cudaMalloc
#define cupmMemcpy                 cudaMemcpy
#define cupmFree                   cudaFree
#define cupmSuccess                cudaSuccess
#define cupmMemcpyHostToDevice     cudaMemcpyHostToDevice
#define PetscMPICUPMAwarenessCheck PetscMPICUDAAwarenessCheck
PetscMPICUPMAwarnessCheckFunction
#endif

#if defined(PETSC_HAVE_HIP)
#define cupmError_t                hipError_t
#define cupmMalloc                 hipMalloc
#define cupmMemcpy                 hipMemcpy
#define cupmFree                   hipFree
#define cupmSuccess                hipSuccess
#define cupmMemcpyHostToDevice     hipMemcpyHostToDevice
#define PetscMPICUPMAwarenessCheck PetscMPIHIPAwarenessCheck
  PetscMPICUPMAwarnessCheckFunction
#endif

#else
void PetscSignalSegvCheckPointerOrMpi(void) {
  return;
}

PetscBool PetscCheckPointer(const void *ptr, PETSC_UNUSED PetscDataType dtype) {
  if (!ptr) return PETSC_FALSE;
  return PETSC_TRUE;
}

#if defined(PETSC_HAVE_CUDA)
PetscBool PetscMPICUDAAwarenessCheck(void) {
  /* If no setjmp (rare), return true and let users code run (and segfault if they should) */
  return PETSC_TRUE;
}
#endif

#if defined(PETSC_HAVE_HIP)
PetscBool PetscMPIHIPAwarenessCheck(void) {
  /* If no setjmp (rare), return true and let users code run (and segfault if they should) */
  return PETSC_TRUE;
}
#endif

#endif