#!/usr/bin/env python
# -*- coding: utf-8 -*-

import click
import wowp.schedulers
import wowp.util
from wowp.logger import logger
import logging
import six
import sys
try:
    from mpi4py import MPI
    IS_MPI = True
    mpi_comm = MPI.COMM_WORLD  # get MPI communicator object
    mpi_size = mpi_comm.size   # total number of processes
    mpi_rank = mpi_comm.rank   # rank of this process
    mpi_status = MPI.Status()  # get MPI status object
except ImportError:
    IS_MPI = False


@click.command(help="Execute a WOW:-P workflow")
@click.option(
    '--scheduler',
    '-s',
    help="Scheduler incl. parameters, default 'LinearizedScheduler()'",
    type=str,
    default="LinearizedScheduler()")
@click.option(
    '--inputs',
    '-i',
    help="File with inputs",
    type=click.Path(exists=True))
@click.option(
    '--output',
    '-o',
    help="Output file, results are saved using wowp.util.dump",
    type=click.Path(exists=False))
@click.option(
    '--arg',
    '-a',
    multiple=True,
    help=("Input port names and values, e.g. -i x 2 will set input port x to 2."
          " Values are processed by eval. Overrides inputs."),
    type=(str, str))
@click.option(
    '--ipykernel/--no-ipykernel',
    help="Start an IPython kernel instead of exiting",
    default=False)
@click.option(
    '--debug/--no-debug',
    help="Switch debug output on",
    default=False)
@click.argument('workflow', nargs=1, required=False, type=click.Path(exists=True))
@click.pass_context
def main(ctx, scheduler, workflow, inputs, output, arg, ipykernel, debug):

    if debug:
        logger.setLevel(logging.DEBUG)

    # detect MPI scheduler
    # TODO this test might be too simple
    if 'mpi' in scheduler.lower() and IS_MPI:
        if mpi_rank > 0:
            print('MPI worker rank {}'.format(mpi_rank))
            from wowp.schedulers import mpi_worker
            mpi_worker()
            return
        else:
            print('This is MPI master :-P')

    print("WOW:-Ping {workflow} with {scheduler}".format(scheduler=scheduler, workflow=workflow))
    with open(workflow, 'r') as wffile:
        try:
            exec_context = {}
            six.exec_(wffile.read(), exec_context)
        except Exception:
            print('Error executing the workflow file:')
            six.reraise(*sys.exc_info())
    if 'WORKFLOW' not in exec_context:
        raise ValueError('WORKFLOW must be defined in the workflow file')
    workflow = exec_context['WORKFLOW']
    # TODO temporary
    # inputs = exec_context['INPUTS']
    kwargs = {}
    if inputs:
        with open(inputs, 'r') as inpfile:
            try:
                exec_context = {}
                six.exec_(inpfile.read(), exec_context)
            except Exception:
                print('Error executing the workflow file:')
                six.reraise(*sys.exc_info())
        if 'INPUTS' not in exec_context or not isinstance(exec_context['INPUTS'], dict):
            raise ValueError('INPUTS dictionary must be defined in the workflow file')
        kwargs.update(exec_context['INPUTS'])
    # assing inputs passed from command line
    for k, v in arg:
        kwargs[k] = eval(v, {})
    print(kwargs)

    try:
        print('init scheduler')
        scheduler = eval("wowp.schedulers.{}".format(scheduler))
    except AttributeError:
        raise('Scheduler {scheduler} not found'.format(scheduler=scheduler))

    # run the workflow
    print("scheduler.run_workflow")
    scheduler.run_workflow(workflow, **kwargs)
    res = {port.name: port.pop_all() for port in workflow.outports}
    print('Result:')
    print(res)
    # save to file
    if output:
        with open(output, 'wb') as fo:
            wowp.util.dump(res, fo)

    if ipykernel:
        from IPython import embed_kernel
        ipy_context = {'workflow': workflow, 'scheduler': scheduler, 'kwargs': kwargs}
        embed_kernel(local_ns=ipy_context)

    scheduler.shutdown()


if __name__ == '__main__':
    main()
