/* StarPU --- Runtime system for heterogeneous multicore architectures.
 *
 * Copyright (C) 2013-2020  Université de Bordeaux, CNRS (LaBRI UMR 5800), Inria
 *
 * StarPU is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation; either version 2.1 of the License, or (at
 * your option) any later version.
 *
 * StarPU is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 *
 * See the GNU Lesser General Public License in COPYING.LGPL for more details.
 */


#include <stdarg.h>
#include <stdlib.h>
#include <common/utils.h>

#include <mpi_failure_tolerance/starpu_mpi_checkpoint.h>
#include <mpi_failure_tolerance/starpu_mpi_checkpoint_template.h>
#include <mpi_failure_tolerance/starpu_mpi_checkpoint_package.h>
#include <sys/param.h>
#include <starpu_mpi_private.h>
#include <mpi/starpu_mpi_mpi_backend.h> // Should be deduced at preprocessing (Nmad vs MPI)
#include <mpi/starpu_mpi_mpi.h>
#include "starpu_mpi_cache.h"


static struct _starpu_mpi_req_list detached_ft_service_requests;
static unsigned detached_send_n_ft_service_requests;
static starpu_pthread_mutex_t detached_ft_service_requests_mutex;

#ifdef STARPU_MPI_VERBOSE
static char *_starpu_mpi_request_type(enum _starpu_mpi_request_type request_type)
{
	switch (request_type)
	{
		case SEND_REQ: return "SEND_REQ";
		case RECV_REQ: return "RECV_REQ";
		case WAIT_REQ: return "WAIT_REQ";
		case TEST_REQ: return "TEST_REQ";
		case BARRIER_REQ: return "BARRIER_REQ";
		case UNKNOWN_REQ: return "UNSET_REQ";
		default: return "unknown request type";
	}
}
#endif

int _ft_service_msg_recv_send_common(void* ptr, int count, int rank, int tag, int req_type, MPI_Comm comm, void (*callback)(void *), void* arg)
{
	struct _starpu_mpi_req* req;

	/* Check if the tag is a service message */
	STARPU_ASSERT_MSG(tag==_STARPU_MPI_TAG_CP_ACK || tag == _STARPU_MPI_TAG_CP_INFO, "Only _STARPU_MPI_TAG_CP_ACK or _STARPU_MPI_TAG_CP_INFO are service msgs.");

	/* Initialize the request structure */
	_starpu_mpi_request_init(&req);
	req->request_type = req_type;
//	/* prio_list is sorted by increasing values */
//	//TODO: Check compatibility with prio
//	if (_starpu_mpi_use_prio)
//		req->prio = 0;
	req->data_handle = NULL;
	req->node_tag.node.rank = rank;
	req->node_tag.data_tag = tag;
	req->node_tag.node.comm = comm;
	req->detached = 1;
	req->ptr = ptr;
	req->sync = 0;
	req->datatype = MPI_BYTE;
	req->callback = callback;
	req->callback_arg = arg;
	req->func = NULL;
	req->sequential_consistency = 1;
	req->count = count;

	_mpi_backend._starpu_mpi_backend_request_fill(req, comm, 0);

	STARPU_PTHREAD_MUTEX_LOCK(&detached_ft_service_requests_mutex);
	if (req_type==SEND_REQ) {
		MPI_Isend(req->ptr, req->count, req->datatype, req->node_tag.node.rank, req->node_tag.data_tag,
		          req->node_tag.node.comm, &req->backend->data_request);
	}
	else if (req_type==RECV_REQ) {
		MPI_Irecv(req->ptr, req->count, req->datatype, req->node_tag.node.rank, req->node_tag.data_tag,
		          req->node_tag.node.comm, &req->backend->data_request);
	}
	else {
		STARPU_ASSERT_MSG(1, "Unrecognized req type: Only RECV_REQ and SEND_REQ are accepeted\n");
	}
	_starpu_mpi_req_list_push_back(&detached_ft_service_requests, req);
	_STARPU_MPI_DEBUG(2, "pushed service req: %p in list %p - prev: %p - next: %p - dest:%d - tag:%d - type:%s\n", req, &detached_ft_service_requests, _starpu_mpi_req_list_prev(req), _starpu_mpi_req_list_next(req), req->node_tag.node.rank, (int)req->node_tag.data_tag, req_type ? "recv" : "send");
	if (req_type==SEND_REQ) {
		detached_send_n_ft_service_requests++;
	}
	req->submitted = 1;

	STARPU_PTHREAD_MUTEX_UNLOCK(&detached_ft_service_requests_mutex);

	_starpu_mpi_wake_up_progress_thread();

	return 0;
}

int _ft_service_msg_isend_cb(void* msg, int count, int rank, int tag, MPI_Comm comm, void (*callback)(void *), void* arg)
{
	return _ft_service_msg_recv_send_common(msg, count, rank, tag, SEND_REQ, comm, callback, arg);
}

int _ft_service_msg_irecv_cb(void* msg, int count, int rank, int tag, MPI_Comm comm, void (*callback)(void *), void* arg)
{
	return _ft_service_msg_recv_send_common(msg, count, rank, tag, RECV_REQ, comm, callback, arg);
}


static void _starpu_mpi_handle_ft_request_termination(struct _starpu_mpi_req *req)
{
	_STARPU_MPI_LOG_IN();
	_STARPU_MPI_DEBUG(2, "complete MPI request %p type %s tag %"PRIi64" src %d data %p ptr %p datatype '%s' count %d registered_datatype %d internal_req %p\n",
			req, _starpu_mpi_request_type(req->request_type), req->node_tag.data_tag, req->node_tag.node.rank, req->data_handle, req->ptr,
			req->datatype_name, (int)req->count, req->registered_datatype, req->backend->internal_req);

	if (req->backend->internal_req)
	{
		free(req->backend->early_data_handle);
		req->backend->early_data_handle = NULL;
	}
	else
	{
		if (req->request_type == RECV_REQ || req->request_type == SEND_REQ)
		{
			if (req->registered_datatype == 0)
			{
				if (req->request_type == SEND_REQ)
				{
					// We need to make sure the communication for sending the size
					// has completed, as MPI can re-order messages, let's call
					// MPI_Wait to make sure data have been sent
					starpu_free_on_node_flags(STARPU_MAIN_RAM, (uintptr_t)req->ptr, req->count, 0);
					req->ptr = NULL;
				}
				else if (req->request_type == RECV_REQ)
				{
					// req->ptr is freed by starpu_data_unpack
					starpu_data_unpack(req->data_handle, req->ptr, req->count);
					starpu_memory_deallocate(STARPU_MAIN_RAM, req->count);
				}
			}
			else
			{
				//_starpu_mpi_datatype_free(req->data_handle, &req->datatype);
			}
		}
		_STARPU_MPI_TRACE_TERMINATED(req, req->node_tag.node.rank, req->node_tag.data_tag);
	}

	_starpu_mpi_release_req_data(req);

	if (req->backend->envelope)
	{
		free(req->backend->envelope);
		req->backend->envelope = NULL;
	}

	/* Execute the specified callback, if any */
	if (req->callback)
		req->callback(req->callback_arg);

	/* tell anyone potentially waiting on the request that it is
	 * terminated now */
	STARPU_PTHREAD_MUTEX_LOCK(&req->backend->req_mutex);
	req->completed = 1;
	STARPU_PTHREAD_COND_BROADCAST(&req->backend->req_cond);
	STARPU_PTHREAD_MUTEX_UNLOCK(&req->backend->req_mutex);
	_STARPU_MPI_LOG_OUT();
}

void starpu_mpi_test_ft_detached_service_requests(void)
{
	//_STARPU_MPI_LOG_IN();
	int flag;
	struct _starpu_mpi_req *req;

	STARPU_PTHREAD_MUTEX_LOCK(&detached_ft_service_requests_mutex);

	if (_starpu_mpi_req_list_empty(&detached_ft_service_requests))
	{
		STARPU_PTHREAD_MUTEX_UNLOCK(&detached_ft_service_requests_mutex);
		//_STARPU_MPI_LOG_OUT();
		return;
	}

	_STARPU_MPI_TRACE_TESTING_DETACHED_BEGIN();
	req = _starpu_mpi_req_list_begin(&detached_ft_service_requests);
	while (req != _starpu_mpi_req_list_end(&detached_ft_service_requests))
	{
		STARPU_PTHREAD_MUTEX_UNLOCK(&detached_ft_service_requests_mutex);

		_STARPU_MPI_TRACE_TEST_BEGIN(req->node_tag.node.rank, req->node_tag.data_tag);
		//_STARPU_MPI_DEBUG(3, "Test detached request %p - mpitag %"PRIi64" - TYPE %s %d\n", &req->backend->data_request, req->node_tag.data_tag, _starpu_mpi_request_type(req->request_type), req->node_tag.node.rank);
#ifdef STARPU_SIMGRID
		req->ret = _starpu_mpi_simgrid_mpi_test(&req->done, &flag);
#else
		STARPU_MPI_ASSERT_MSG(req->backend->data_request != MPI_REQUEST_NULL, "Cannot test completion of the request MPI_REQUEST_NULL");
		req->ret = MPI_Test(&req->backend->data_request, &flag, MPI_STATUS_IGNORE);
#endif

		STARPU_MPI_ASSERT_MSG(req->ret == MPI_SUCCESS, "MPI_Test returning %s", _starpu_mpi_get_mpi_error_code(req->ret));
		_STARPU_MPI_TRACE_TEST_END(req->node_tag.node.rank, req->node_tag.data_tag);

		if (!flag)
		{
			req = _starpu_mpi_req_list_next(req);
		}
		else
		{
			_STARPU_MPI_TRACE_POLLING_END();
			struct _starpu_mpi_req *next_req;
			next_req = _starpu_mpi_req_list_next(req);

			_STARPU_MPI_TRACE_COMPLETE_BEGIN(req->request_type, req->node_tag.node.rank, req->node_tag.data_tag);

			STARPU_PTHREAD_MUTEX_LOCK(&detached_ft_service_requests_mutex);
			if (req->request_type == SEND_REQ)
				detached_send_n_ft_service_requests--;
			_starpu_mpi_req_list_erase(&detached_ft_service_requests, req);
			STARPU_PTHREAD_MUTEX_UNLOCK(&detached_ft_service_requests_mutex);
			_starpu_mpi_handle_ft_request_termination(req);

			_STARPU_MPI_TRACE_COMPLETE_END(req->request_type, req->node_tag.node.rank, req->node_tag.data_tag);

			STARPU_PTHREAD_MUTEX_LOCK(&req->backend->req_mutex);
			/* We don't want to free internal non-detached
			   requests, we need to get their MPI request before
			   destroying them */
			if (req->backend->is_internal_req && !req->backend->to_destroy)
			{
				/* We have completed the request, let the application request destroy it */
				req->backend->to_destroy = 1;
				STARPU_PTHREAD_MUTEX_UNLOCK(&req->backend->req_mutex);
			}
			else
			{
				STARPU_PTHREAD_MUTEX_UNLOCK(&req->backend->req_mutex);
				_starpu_mpi_request_destroy(req);
			}

			req = next_req;
			_STARPU_MPI_TRACE_POLLING_BEGIN();
		}

		STARPU_PTHREAD_MUTEX_LOCK(&detached_ft_service_requests_mutex);
	}
	_STARPU_MPI_TRACE_TESTING_DETACHED_END();

	STARPU_PTHREAD_MUTEX_UNLOCK(&detached_ft_service_requests_mutex);
	//_STARPU_MPI_LOG_OUT();
}

int starpu_mpi_ft_service_lib_init()
{
	_starpu_mpi_req_list_init(&detached_ft_service_requests);
	starpu_pthread_mutex_init(&detached_ft_service_requests_mutex, NULL);

	return 0;
}

int starpu_mpi_ft_service_lib_busy()
{
	return !_starpu_mpi_req_list_empty(&detached_ft_service_requests);
}