#!/usr/bin/python

# Copyright (c) 2011 Christian Haselgrove
# Licensed under the BSD license: http://www.opensource.org/licenses/BSD-2-Clause

import sys
import os
import getopt
import getpass
import datetime
import urllib2
import xml.dom.minidom
import suds.client
import suds.xsd.doctor
import suds.transport.http

class XNATSoapCaller:

    def __init__(self, username, password):
        self.cookiejar = None
        result = self.call('CreateServiceSession.jws', 
                           'execute', 
                           (), 
                           auth=(username, password))
        self.session = str(result)
        return

    def close(self):
        self.call('CloseServiceSession.jws', 'execute', ())
        return

    def call(self, jws, operation, inputs, fix_import=False, auth=None):
        url = '%s/axis/%s' % ('http://candi-store/xnat/', jws)
        if auth:
            t = suds.transport.http.HttpAuthenticated(username=auth[0], 
                                                      password=auth[1])
        else:
            t = suds.transport.http.HttpTransport()
        if self.cookiejar:
            t.cookiejar = self.cookiejar
        if fix_import:
            xsd_url = 'http://schemas.xmlsoap.org/soap/encoding/'
            imp = suds.xsd.doctor.Import(xsd_url)
            doctor = suds.xsd.doctor.ImportDoctor(imp)
            client = suds.client.Client('%s?wsdl' % url, 
                                        transport=t, 
                                        doctor=doctor)
        else:
            client = suds.client.Client('%s?wsdl' % url, transport=t)
        typed_inputs = []
        for (dtype, val) in inputs:
            ti = client.factory.create(dtype)
            ti.value = val
            typed_inputs.append(ti)
        # the WSDL returns the local IP address in the URLs; these need 
        # to be corrected if XNAT is behind a proxy
        client.set_options(location=url)
        f = getattr(client.service, operation)
        result = f(*typed_inputs)
        if not self.cookiejar:
            self.cookiejar = t.cookiejar
        return result

def report_error(msg):
	sys.stderr.write('%s: %s\n' % (progname, msg))
	sys.stderr.write('run %s with no arguments for usage\n' % progname)
	return

def node_text(node):
	s = ''
	for cn in node.childNodes:
		if cn.nodeType == cn.TEXT_NODE:
			s += cn.data
	return s

def print_pairs(pairs):
	if not pairs:
		return
	max_name_length = max([ len(name) for (name, value) in pairs ])
	fmt = '%%-%ds %%s' % max_name_length
	for pair in pairs:
		print fmt % pair
			
def get_workflow_xml(xsc, experiment_id, w_id):
	args = (('ns1:string', xsc.session), 
        	('ns1:string', 'wrk:workflowData.ID'), 
        	('ns1:string', '='), 
        	('ns1:string', experiment_id), 
        	('ns1:string', 'wrk:workflowData'))
	w_ids = [ int(v) for v in xsc.call('GetIdentifiers.jws', 'search', args) ]
	if w_id not in w_ids:
		sys.stderr.write("%s: can't find ID %d for %s (maybe try '%s list %s')\n" % (progname, w_id, experiment_id, progname, experiment_id))
		sys.exit(1)
	url = '%s/app/template/XMLSearch.vm/id/%s/data_type/wrk:workflowData' % (host, str(w_id))
	r = urllib2.Request(url)
	xsc.cookiejar.add_cookie_header(r)
	return urllib2.urlopen(r).read()

def update(xsc, doc):
	inputs = (('ns0:string', xsc.session), 
	          ('ns0:string', doc.toxml()), 
	          ('ns0:boolean', False), 
	          ('ns0:boolean', True))
	xsc.call('StoreXML.jws', 'store', inputs, fix_import=True)
	return

progname = os.path.basename(sys.argv[0])

try:
	host = os.environ['XNAT_HOST']
except KeyError:
	host = None

try:
	user_name = os.environ['XNAT_USER']
except KeyError:
	user_name = None

try:
	password = os.environ['XNAT_PASSWORD']
except KeyError:
	password = None

if len(sys.argv) == 1:
	print
	print 'usage: %s [options] <command> [command arguments]' % progname
	print
	print 'XNAT workflow tool'
	print
	print 'options are:'
	print
	print '    -h <host>'
	print '    -u <user name>'
	print '    -p <password>'
	print
	print 'option values may be given by environment variables:'
	print
	if host is None:
		print '    host: XNAT_HOST'
	else:
		print '    host: XNAT_HOST (set to %s)' % host
	if user_name is None:
		print '    user name: XNAT_USER'
	else:
		print '    user name: XNAT_USER (set to %s)' % user_name
	if password is None:
		print '    password: XNAT_PASSWORD'
	else:
		print '    password: XNAT_PASSWORD (currently set)'
	print
	print '%s will prompt for a missing user name, password, or host' % progname
	print
	print 'commands are:'
	print
	print '    list <experiment identifier>'
	print '    info <experiment identifier> <workflow ID>'
	print '    xml <experiment identifier> <workflow ID>'
	print '    update <experiment identifier> <workflow ID> <step ID> <step description> <percentage complete>'
	print '    complete <experiment identifier> <workflow ID>'
	print '    fail <experiment identifier> <workflow ID> [step description]'
	print
	sys.exit(1)

try:
	(opts, args) = getopt.getopt(sys.argv[1:], 'h:u:p:')
except getopt.error, data:
	report_error(data)
	sys.exit(1)

if not args:
	report_error('not enough non-option arguments')
	sys.exit(1)

for (option, value) in opts:
	if option == '-h':
		host = value
	if option == '-u':
		user_name = value
	if option == '-p':
		password = value

if host is None:
	sys.stdout.write('Host: ')
	sys.stdout.flush()
	host = sys.stdin.readline().strip()

if user_name is None:
	sys.stdout.write('User name: ')
	sys.stdout.flush()
	user_name = sys.stdin.readline().strip()

if password is None:
	password = getpass.getpass()

host = host.rstrip('/')

command = args.pop(0)

if command == 'list':
	if not args:
		report_error('"list" requires an argument')
		sys.exit(0)
	experiment_id = args[0]
	xsc = XNATSoapCaller(user_name, password)
	args = (('ns1:string', xsc.session), 
        	('ns1:string', 'wrk:workflowData.ID'), 
        	('ns1:string', '='), 
        	('ns1:string', experiment_id), 
        	('ns1:string', 'wrk:workflowData'))
	for w_id in xsc.call('GetIdentifiers.jws', 'search', args):
		url = '%s/app/template/XMLSearch.vm/id/%s/data_type/wrk:workflowData' % (host, str(w_id))
		r = urllib2.Request(url)
		xsc.cookiejar.add_cookie_header(r)
		data = urllib2.urlopen(r).read()
		doc = xml.dom.minidom.parseString(data)
		workflow_node = doc.getElementsByTagName('wrk:Workflow')[0]
		print w_id, \
		      workflow_node.getAttribute('ExternalID'), \
		      workflow_node.getAttribute('ID'), \
		      workflow_node.getAttribute('status'), \
		      workflow_node.getAttribute('pipeline_name'), \
		      workflow_node.getAttribute('launch_time'), \
		      workflow_node.getAttribute('percentageComplete')
elif command == 'xml':
	if len(args) != 2:
		report_error('"xml" requires two arguments')
		sys.exit(0)
	experiment_id = args[0]
	try:
		w_id = int(args[1])
	except ValueError:
		report_error('bad workflow ID "%s"' % args[1])
		sys.exit(1)
	xsc = XNATSoapCaller(user_name, password)
	print get_workflow_xml(xsc, experiment_id, w_id)
elif command == 'info':
	if len(args) != 2:
		report_error('"info" requires two arguments')
		sys.exit(0)
	experiment_id = args[0]
	try:
		w_id = int(args[1])
	except ValueError:
		report_error('bad workflow ID "%s"' % args[1])
		sys.exit(1)
	xsc = XNATSoapCaller(user_name, password)
	doc = xml.dom.minidom.parseString(get_workflow_xml(xsc, experiment_id, w_id))
	workflow_node = doc.getElementsByTagName('wrk:Workflow')[0]
	print '--- basic info'
	print
	for attr in ('data_type', 
	             'ID', 
	             'ExternalID', 
	             'status', 
	             'pipeline_name', 
	             'step_description', 
	             'launch_time', 
	             'current_step_id', 
	             'current_step_launch_time', 
	             'percentageComplete'):
			print '%-25s %s' % (attr + ':', workflow_node.getAttribute(attr))
	print
	environment_nodes = doc.getElementsByTagName('wrk:Workflow')
	if environment_nodes:
		print '--- execution environment'
		print
		pairs = []
		for tag in ('pipeline', 
		            'xnatuser', 
		            'host', 
		            'notify', 
		            'dataType', 
		            'id', 
		            'supressNotification'):
			elements = environment_nodes[0].getElementsByTagName('wrk:' + tag)
			if not elements:
				continue
			pairs.append((tag+':', node_text(elements[0])))
		print_pairs(pairs)
		print
		print '--- parameters'
		print
		pairs = []
		for pn in environment_nodes[0].getElementsByTagName('wrk:parameter'):
			name = pn.getAttribute('name')
			value = node_text(pn)
			pairs.append((name+':', value))
		print_pairs(pairs)
		print
	xsc.close()
elif command == 'update':
	if len(args) != 5:
		report_error('"update" requires five arguments')
		sys.exit(0)
	experiment_id = args[0]
	try:
		w_id = int(args[1])
	except ValueError:
		report_error('bad workflow ID "%s"' % args[1])
		sys.exit(1)
	xsc = XNATSoapCaller(user_name, password)
	doc = xml.dom.minidom.parseString(get_workflow_xml(xsc, experiment_id, w_id))
	workflow_node = doc.getElementsByTagName('wrk:Workflow')[0]
	workflow_node.setAttribute('status', 'Running')
	t = datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')
	workflow_node.setAttribute('current_step_launch_time', t)
	workflow_node.setAttribute('current_step_id', args[2])
	workflow_node.setAttribute('step_description', args[3])
	workflow_node.setAttribute('percentageComplete', args[4])
	update(xsc, doc)
	xsc.close()
	print 'updated workflow %d for %s' % (w_id, experiment_id)
elif command == 'complete':
	if len(args) != 2:
		report_error('"complete" requires two arguments')
		sys.exit(0)
	experiment_id = args[0]
	try:
		w_id = int(args[1])
	except ValueError:
		report_error('bad workflow ID "%s"' % args[1])
		sys.exit(1)
	xsc = XNATSoapCaller(user_name, password)
	doc = xml.dom.minidom.parseString(get_workflow_xml(xsc, experiment_id, w_id))
	workflow_node = doc.getElementsByTagName('wrk:Workflow')[0]
	workflow_node.setAttribute('status', 'Complete')
	t = datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')
	workflow_node.setAttribute('current_step_launch_time', t)
	workflow_node.setAttribute('percentageComplete', '100.0')
	try:
		workflow_node.removeAttribute('current_step_id')
	except xml.dom.NotFoundErr:
		pass
	try:
		workflow_node.removeAttribute('step_description')
	except xml.dom.NotFoundErr:
		pass
	update(xsc, doc)
	xsc.close()
	print 'closed (completed) workflow %d for %s' % (w_id, experiment_id)
elif command == 'fail':
	if len(args) < 2 or len(args) > 3:
		report_error('"fail" requires two or three arguments')
		sys.exit(0)
	experiment_id = args[0]
	try:
		w_id = int(args[1])
	except ValueError:
		report_error('bad workflow ID "%s"' % args[1])
		sys.exit(1)
	xsc = XNATSoapCaller(user_name, password)
	doc = xml.dom.minidom.parseString(get_workflow_xml(xsc, experiment_id, w_id))
	workflow_node = doc.getElementsByTagName('wrk:Workflow')[0]
	workflow_node.setAttribute('status', 'Failed')
	if len(args) > 2:
		workflow_node.setAttribute('step_description', args[2])
	update(xsc, doc)
	xsc.close()
	print 'closed (failed) workflow %d for %s' % (w_id, experiment_id)
else:
	report_error('unknown command "%s"' % command)
	sys.exit(1)

sys.exit(0)

# eof