#!/usr/bin/env python
import base64
import os

from torii import Torii
from torii.data import Task

ICON = ''

def picom_body(ex):
    """
    The function defining the content of the Picom
    :param ex: the picom executor
    :return:
    """

    sleep = 10
    nb_taches = 5

    if ex.context.step == 0:
        tr_1 = ex.tasks.new(template='swot tr1', name='tr_1', submit=True)
        tr_1.set_input('sleep', sleep)
        tr_1.context.cluster = ex.context.cluster

        tr_2 = ex.tasks.new(template='swot tr2', name='tr_2', submit=True)
        tr_2.add_dependency(tr_1)
        tr_2.set_input('nb_taches', nb_taches)
        tr_2.map_input(name='sleep', task=tr_1, input='sleep')
        tr_2.context.cluster = ex.context.cluster

        tr_3 = ex.tasks.new(template='swot tr3', name='tr_3', submit=True)
        tr_4 = ex.tasks.new(template='swot tr4', name='tr_4', submit=True)
        tr_3.add_dependency(tr_2)
        tr_4.add_dependency(tr_2)
        tr_3.map_input(name='sleep', task=tr_2, input='sleep')
        tr_4.map_input(name='sleep', task=tr_2, input='sleep')
        tr_3.context.cluster = ex.context.cluster
        tr_4.context.cluster = ex.context.cluster

    else:
        status = [task.status for task in ex.last_tasks]
        if 'CANCELLED' in status or 'FAILED' in status:
            ex.set_status(status='FAILED', message='The workflow has failed')
        else:
            ex.set_status(status='SUCCESS', message='The workflow has succeed')

    return




t = Torii(container='new')

def cnes_swot_example_1():

    print('step 0 - open the picom to edit')
    name = 'swot'
    version = '0.1'
    picom = t.picoms.get(name=name, version=version, take_first=True)
    if not picom:
        # ... or create it
        picom = t.picoms.new(name=name, version=version)

    print('step 1 - set the script from a function')
    picom.set_script(picom_body)

    print('step 2 - import all the applications')
    application_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'applications', 'swot_*')
    applications = t.applications.import_application(path=application_dir, force=True)

    print('step 3 - prepare the task templates for the picom')
    picom.templateTasks = []
    for application in applications:
        picom.templateTasks.append(t.tasks.new(name = application.name.lower(),
                                               description = application.description,
                                               application=application))

    print('step 4 - update the picom')
    # define the icon
    icon_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "icon.png")
    with open(icon_path, "rb") as file:
        picom.icon = 'data:image/png;base64,{0}'.format(base64.b64encode(file.read()))

    # save the picom
    picom.update()

    print('step 5 - create a new picom task and submit')
    picom_task = t.tasks.new(name = 'operator task', description = 'start a sismic threads and operate them', picom=picom)
    cluster = t.clusters.get(name='devil')
    picom_task.set_cluster(cluster)
    picom_task.submit()

    print('step 1 - wait for the completion')
    picom_task.wait(status=['PAUSED'] + Task.COMPLETE_STATUS)

    print('Task {0} is {1}'.format(picom_task, picom_task.status))
    assert picom_task.status == 'FINISHED'


if __name__ == '__main__':
    cnes_swot_example_1()
