/*****************************************************************************\
 *  step_ctx.c - step_ctx task functions for use by AIX/POE
 *
 *  $Id$
 *****************************************************************************
 *  Copyright (C) 2004 The Regents of the University of California.
 *  Produced at Lawrence Livermore National Laboratory (cf, DISCLAIMER).
 *  Written by Morris Jette <jette1@llnl.gov>.
 *  UCRL-CODE-217948.
 *  
 *  This file is part of SLURM, a resource management program.
 *  For details, see <http://www.llnl.gov/linux/slurm/>.
 *  
 *  SLURM is free software; you can redistribute it and/or modify it under
 *  the terms of the GNU General Public License as published by the Free
 *  Software Foundation; either version 2 of the License, or (at your option)
 *  any later version.
 *  
 *  SLURM 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 General Public License for more
 *  details.
 *  
 *  You should have received a copy of the GNU General Public License along
 *  with SLURM; if not, write to the Free Software Foundation, Inc.,
 *  59 Temple Place, Suite 330, Boston, MA  02111-1307  USA.
\*****************************************************************************/

#ifdef HAVE_CONFIG_H
#  include "config.h"
#endif

#include <errno.h>
#include <pthread.h>
#include <stdarg.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <netinet/in.h>
#include <sys/param.h>
#include <sys/socket.h>
#include <sys/types.h>

#include <slurm/slurm.h>

#include "src/common/hostlist.h"
#include "src/common/net.h"
#include "src/common/slurm_protocol_api.h"
#include "src/common/slurm_protocol_defs.h"
#include "src/common/xmalloc.h"
#include "src/common/xstring.h"

#include "src/api/step_ctx.h"

static void	_xcopy_char_array(char ***argv_p, char **argv, int cnt);
static void	_xfree_char_array(char ***argv_p, int cnt);
static job_step_create_request_msg_t *_copy_step_req(
	job_step_create_request_msg_t *step_req);

/*
 * slurm_step_ctx_create - Create a job step and its context. 
 * IN step_req - description of job step request
 * RET the step context or NULL on failure with slurm errno set
 * NOTE: Free allocated memory using slurm_step_ctx_destroy.
 */
extern slurm_step_ctx
slurm_step_ctx_create (const job_step_create_request_msg_t *user_step_req)
{
	struct slurm_step_ctx_struct *ctx = NULL;
	job_step_create_request_msg_t *step_req = NULL;
	job_step_create_response_msg_t *step_resp = NULL;
	int sock = -1;
	short port = 0;
	int errnum = 0;
	
	/* First copy the user's user_step_req struct in case we
	 * need to modify anything before contacting the slurmctld
	 */
	step_req =
		_copy_step_req((job_step_create_request_msg_t *)user_step_req);

	/* If step_req->host is NULL, then the user wants us
	 * to handle messages from the controller on our own.  We handle
	 * the messages in the step_launch.c mesage handler, but we
	 * need to open the socket right now so we can tell the
	 * controller which port to use.
	 */
	if (step_req->host == NULL) {
		if (net_stream_listen(&sock, &port) < 0) {
			errnum = errno;
			error("unable to intialize step context socket: %m");
			slurm_free_job_step_create_request_msg(step_req);
			goto fail;
		}
		step_req->port = port;
		step_req->host = xshort_hostname();
	}

	if ((slurm_job_step_create(step_req, &step_resp) < 0) ||
	    (step_resp == NULL)) {
		errnum = errno;
		slurm_free_job_step_create_request_msg(step_req);
		goto fail;
	}
	
	ctx = xmalloc(sizeof(struct slurm_step_ctx_struct));
	ctx->launch_state = NULL;
	ctx->magic	= STEP_CTX_MAGIC;
	ctx->job_id	= step_req->job_id;
	ctx->user_id	= step_req->user_id;
	ctx->step_req   = step_req; /* is already a copy of user_step_req */
	ctx->step_resp	= step_resp;

	ctx->launch_state = step_launch_state_create(ctx);
	ctx->launch_state->slurmctld_socket_fd = sock;

fail:
	errno = errnum;
	return (slurm_step_ctx)ctx;
}

/*
 * slurm_step_ctx_get - get parameters from a job step context.
 * IN ctx - job step context generated by slurm_step_ctx_create
 * RET SLURM_SUCCESS or SLURM_ERROR (with slurm_errno set)
 */
extern int
slurm_step_ctx_get (slurm_step_ctx ctx, int ctx_key, ...)
{
	va_list ap;
	int rc = SLURM_SUCCESS;
	uint32_t node_inx;
	uint32_t *uint32_ptr;
	uint32_t **uint32_array_pptr = (uint32_t **) NULL;
	char **char_array_pptr = (char **) NULL;
	job_step_create_response_msg_t ** step_resp_pptr;
	slurm_cred_t  *cred;     /* Slurm job credential    */
	switch_jobinfo_t *switch_job;
	
	if ((ctx == NULL) ||
	    (ctx->magic != STEP_CTX_MAGIC)) {
		slurm_seterrno(EINVAL);
		return SLURM_ERROR;
	}

	va_start(ap, ctx_key);
	switch (ctx_key) {
	case SLURM_STEP_CTX_JOBID:
		uint32_ptr = (uint32_t *) va_arg(ap, void *);
		*uint32_ptr = ctx->job_id;
		break;
	case SLURM_STEP_CTX_STEPID:
		uint32_ptr = (uint32_t *) va_arg(ap, void *);
		*uint32_ptr = ctx->step_resp->job_step_id;
		break;
	case SLURM_STEP_CTX_TASKS:
		uint32_array_pptr = (uint32_t **) va_arg(ap, void *);
		*uint32_array_pptr = ctx->step_resp->step_layout->tasks;
		break;
		
	case SLURM_STEP_CTX_TID:
		node_inx = va_arg(ap, uint32_t);
		if ((node_inx < 0)
		    || (node_inx > ctx->step_resp->step_layout->node_cnt)) {
			slurm_seterrno(EINVAL);
			rc = SLURM_ERROR;
			break;
		}
		uint32_array_pptr = (uint32_t **) va_arg(ap, void *);
		*uint32_array_pptr =
			ctx->step_resp->step_layout->tids[node_inx];
		break;
		
	case SLURM_STEP_CTX_RESP:
		step_resp_pptr = (job_step_create_response_msg_t **) 
			va_arg(ap, void *);
		*step_resp_pptr = ctx->step_resp;
		break;
	case SLURM_STEP_CTX_CRED:
		cred = (slurm_cred_t *) va_arg(ap, void *);
		*cred = ctx->step_resp->cred;
		break;
	case SLURM_STEP_CTX_SWITCH_JOB:
		switch_job = (switch_jobinfo_t *) va_arg(ap, void *);
		*switch_job = ctx->step_resp->switch_job;
		break;
	case SLURM_STEP_CTX_NUM_HOSTS:
		uint32_ptr = (uint32_t *) va_arg(ap, void *);
		*uint32_ptr = ctx->step_resp->step_layout->node_cnt;
		break;
	case SLURM_STEP_CTX_HOST:
		node_inx = va_arg(ap, uint32_t);
		if ((node_inx < 0)
		    || (node_inx > ctx->step_resp->step_layout->node_cnt)) {
			slurm_seterrno(EINVAL);
			rc = SLURM_ERROR;
			break;
		}
		char_array_pptr = (char **) va_arg(ap, void *);
		*char_array_pptr = nodelist_nth_host(
			ctx->step_resp->step_layout->node_list, node_inx);
		break;
	default:
		slurm_seterrno(EINVAL);
		rc = SLURM_ERROR;
	}
	va_end(ap);

	return rc;
}

/*
 * slurm_jobinfo_ctx_get - get parameters from jobinfo context.
 * IN jobinfo - job information from context, returned by slurm_step_ctx_get()
 * IN data_type - type of data required, specific to the switch type
 * OUT data - the requested data type
 * RET SLURM_SUCCESS or SLURM_ERROR (with slurm_errno set)
 */
extern int
slurm_jobinfo_ctx_get(switch_jobinfo_t jobinfo, int data_type, void *data)
{
	if (jobinfo == NULL) {
		slurm_seterrno(EINVAL);
		return SLURM_ERROR;
	}

	return switch_g_get_jobinfo(jobinfo, data_type, data);
}

/*
 * slurm_step_ctx_set - set parameters in job step context.
 * IN ctx - job step context generated by slurm_step_ctx_create
 * RET SLURM_SUCCESS or SLURM_ERROR (with slurm_errno set)
 */
extern int
slurm_step_ctx_set (slurm_step_ctx ctx, int ctx_key, ...)
{
	va_list ap;
	int rc = SLURM_SUCCESS;

	if ((ctx == NULL) ||
	    (ctx->magic != STEP_CTX_MAGIC)) {
		slurm_seterrno(EINVAL);
		return SLURM_ERROR;
	}

	va_start(ap, ctx_key);
	switch (ctx_key) {
	case SLURM_STEP_CTX_ARGS:
		if (ctx->argv)
			_xfree_char_array(&ctx->argv, ctx->argc);
		ctx->argc = va_arg(ap, int);
		if ((ctx->argc < 1) || (ctx->argc > 1024)) {
			slurm_seterrno(EINVAL);
			break;
		}
		_xcopy_char_array(&ctx->argv, va_arg(ap, char **), 
				  ctx->argc);
		break;

	case SLURM_STEP_CTX_CHDIR:
		if (ctx->cwd)
			xfree(ctx->cwd);
		ctx->cwd = xstrdup(va_arg(ap, char *));
		break;

	case SLURM_STEP_CTX_ENV:
		ctx->env_set = 1;
		if (ctx->env)
			_xfree_char_array(&ctx->env, ctx->envc);
		ctx->envc = va_arg(ap, int);
		if ((ctx->envc < 1) || (ctx->envc > 1024)) {
			slurm_seterrno(EINVAL);
			break;
		}
		_xcopy_char_array(&ctx->env, va_arg(ap, char **), 
				  ctx->envc);
		break;
	default:
		slurm_seterrno(EINVAL);
		rc = SLURM_ERROR;
	}
	va_end(ap);

	return rc;
}


/*
 * slurm_step_ctx_destroy - free allocated memory for a job step context.
 * IN ctx - job step context generated by slurm_step_ctx_create
 * RET SLURM_SUCCESS or SLURM_ERROR (with slurm_errno set)
 */
extern int
slurm_step_ctx_destroy (slurm_step_ctx ctx)
{
	if ((ctx == NULL) ||
	    (ctx->magic != STEP_CTX_MAGIC)) {
		slurm_seterrno(EINVAL);
		return SLURM_ERROR;
	}
	slurm_free_job_step_create_request_msg(ctx->step_req);
	slurm_free_job_step_create_response_msg(ctx->step_resp);
	if (ctx->argv)
		_xfree_char_array(&ctx->argv, ctx->argc);
	if (ctx->env_set)
		_xfree_char_array(&ctx->env, ctx->envc);
	xfree(ctx->cwd);
	step_launch_state_destroy(ctx->launch_state);
	xfree(ctx);
	return SLURM_SUCCESS;
}

/* xfree an array of character strings as created by _xcopy_char_array */
static void _xfree_char_array(char ***argv_p, int cnt)
{
	char **argv = *argv_p;
	int i;
	for (i=0; i<cnt; i++)
		xfree(argv[i]);
	xfree(*argv_p);
}

/* copy a character array, free with _xfree_char_array */
static void _xcopy_char_array(char ***argv_p, char **argv, int cnt)
{
	int i;
	char **tmp = xmalloc(sizeof(char *) * cnt);

	for (i=0; i<cnt; i++)
		tmp[i] = xstrdup(argv[i]);

	*argv_p = tmp;
}

static job_step_create_request_msg_t *_copy_step_req(
	job_step_create_request_msg_t *step_req)
{
	job_step_create_request_msg_t *copy;

	copy = xmalloc(sizeof(job_step_create_request_msg_t));
	if (copy == NULL)
		return NULL;

	memcpy(copy, step_req, sizeof(job_step_create_request_msg_t));
	if (step_req->host != NULL)
		copy->host = xstrdup(step_req->host);
	if (step_req->node_list != NULL)
		copy->node_list = xstrdup(step_req->node_list);
	if (step_req->network != NULL)
		copy->network = xstrdup(step_req->network);
	if (step_req->name != NULL)
		copy->name = xstrdup(step_req->name);

	return copy;
}

