#!/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:
    """

    nb_tiles = 20
    auto_start = False

    if ex.context.step == 0:
        l1b = ex.tasks.new(template='L1B_SP', name='L1B_SP', submit=auto_start)
        l1b.set_input('nb_tiles', nb_tiles)
        l1b.set_input('goal', 'OK')
        l1b.context.cluster = ex.context.cluster
        l1b.runInfos.cnes_status = 'READY'

        l2 = ex.tasks.new(template='L2_SP', name='L2_SP', submit=auto_start)
        l2.add_dependency(l1b)
        l2.set_input('nb_tiles', nb_tiles)
        l2.set_input('goal', 'POK')
        l2.context.cluster = ex.context.cluster
        l2.map_input(name='nb_tiles', task=l1b, input='nb_tiles')
        l2.runInfos.cnes_status = 'READY'

        l2_riv_sp = ex.tasks.new(template='L2_RIV_SP', name='L2_RIV_SP', submit=auto_start)
        l2_riv_sp.add_dependency(l2)
        l2_riv_sp.set_input('nb_tiles', nb_tiles)
        l2_riv_sp.set_input('goal', 'KO')
        l2_riv_sp.context.cluster = ex.context.cluster
        l2_riv_sp.map_input(name='nb_tiles', task=l1b, input='nb_tiles')
        l2_riv_sp.runInfos.cnes_status = 'READY'

        l2_lak_sp = ex.tasks.new(template='L2_LAK_SP', name='L2_LAK_SP', submit=auto_start)
        l2_lak_sp.add_dependency(l2)
        l2_lak_sp.set_input('nb_tiles', nb_tiles)
        l2_lak_sp.set_input('goal', 'OK')
        l2_lak_sp.context.cluster = ex.context.cluster
        l2_lak_sp.map_input(name='nb_tiles', task=l1b, input='nb_tiles')
        l2_lak_sp.runInfos.cnes_status = 'READY'


    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='FINISHED', message='The workflow has succeed')
        ex.set_status(status='FINISHED', message='The workflow has succeed')

    return


name = 'KaRIn HR'
version = '0.1'


def update_picom(t):

    print('step 0 - open the picom to edit')
    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', '*_SP')
    applications = t.applications.import_application(path=application_dir, force=True)

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

    # templates['L2_SP'].add_dependency(templates['L1B_SP'])
    # templates['PGE_L2_RIVER_TILE'].add_dependency(templates['L2_SP'])
    # templates['PGE_L2_RIV_SP'].add_dependency(templates['PGE_L2_RIVER_TILE'])
    # templates['PGE_L2_LAKE_TILE'].add_dependency(templates['PGE_L2_RIVER_TILE'])
    # templates['PGE_L2_LAK_SP'].add_dependency(templates['PGE_L2_LAKE_TILE'])
    #
    # for template in templates.itervalues():
    #     template.update()

    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()

def run_picom_task(t):

    print('step 0 - get the picom to run')
    picom = t.picoms.get(name=name, version=version, take_first=True)

    print('step 1 - create a new picom task and submit')
    picom_task = t.tasks.new(name = 'KaRIn HR', picom=picom)
    cluster = t.clusters.get(name='devil')
    picom_task.set_cluster(cluster)
    picom_task.submit()

    print('step 2 - 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__':
    t = Torii(container='new', name="my KaRIn HR")

    update_picom(t)
    run_picom_task(t)
