#!/usr/bin/env python
from __future__ import print_function
import os
import sys
import yaml
from xbowflow.pipelines import InterfaceKernel, SubprocessKernel, Pipeline
from xbowflow.clients import dask_client
import argparse

parser = argparse.ArgumentParser(description='Run an xbowflow workflow')
parser.add_argument('configfile', help='Name of the workflow configuration file')
parser.add_argument('inputsfile', help='Name of the workflow job inputs file')

args = parser.parse_args()

kernels = {
           'interfacekernel': InterfaceKernel,
           'subprocesskernel': SubprocessKernel,
          }

with open(args.configfile) as f:
    wconfig = yaml.load(f)

with open(args.inputsfile) as f:
    inputs = yaml.load(f)
# Set the reserved variable CWD to the current working directory:
if not 'CWD' in inputs:
    inputs['CWD'] = os.getcwd()

if 'inputs' in wconfig:
    for k in wconfig['inputs']:
        if not k in inputs:
            raise RuntimeError('Error - inputs is missing parameter {}'.format(k))

if wconfig['workflow'].get('dryrun', False):
    client = None
else:
    client = dask_client()

mykernels = {}
for stage in wconfig['workflow']['stages']:
    for step in wconfig[stage]['steps']:
        ktype = wconfig[step]['type']
        mykernels[step] = kernels[ktype](wconfig[step]['configuration'])

output = inputs
for stage in wconfig['workflow']['stages']:
    klist = [mykernels[k] for k in wconfig[stage]['steps']]
    pipeline = Pipeline(client, klist)
    if 'iterations' in wconfig[stage]:
        n_its = inputs[wconfig[stage]['iterations']]
    else:
        n_its = 1
    for i in range(n_its):
        output = pipeline.run(output)

ok = True
if isinstance(output, list):
    for o in output:
        if o['returncode'] != 0:
            print(o['output'].decode('utf-8'))
            ok = False
else:
    if output['returncode'] != 0:
        print(output['output'].decode('utf-8'))
        ok = False
if ok:
    print('Workflow completed without errors')
