import sys
import inspect

from datacloud.launcher.scenario import slow, fast
from datacloud.launcher.update_cm.app.template import template
from datacloud.dev_utils.logging.logger import get_basic_logger


logger = get_basic_logger(__name__)


IS_CM_TASK = 'is_cm_task'

N_WORKERS_TO_LINE = {
    1: 'MAIN',
    2: 'SCATTER',
    4: 'SCATTER4',
    8: 'SCATTER8',
}


def _generate_declaration(task):
    assert task.is_cm_task

    line = N_WORKERS_TO_LINE[task.cm_data.n_workers]
    line += ' ' + task.cm_data.program_name

    suffix = ''
    if task.cm_data.depends_on:
        suffix += ' ' + ' '.join([
            parent.__name__ for parent in task.cm_data.depends_on
        ])

    if task.cm_data.mailto:
        suffix += ' mailto={}'.format(task.cm_data.mailto)
    if task.cm_data.on_success:
        suffix += ' restart_on_success="{}"'.format(task.cm_data.on_success)
    if task.cm_data.on_failure:
        suffix += ' retry_on_failure="{}"'.format(task.cm_data.on_failure)
    if task.cm_data.timeout:
        suffix += ' timeout="{}"'.format(task.cm_data.timeout)

    if suffix:
        line += ' : ' + suffix
    return line


def _generate_definition(task):
    assert task.is_cm_task
    return '{name}(){{\nrun_cmd "$LAUNCHER --yt-token-name={robot} --program {program}"\n}}'.format(
        name=task.cm_data.program_name,
        robot=task.cm_data.robot,
        program=task.cm_data.program_name
    )


def _build_declaration(task, all_tasks, added, declarations):
    if task.cm_data.program_name in added:
        return
    for parent in task.cm_data.depends_on:
        if parent not in added:
            _build_declaration(parent, all_tasks, added, declarations)

    declarations.append(_generate_declaration(task))
    added.add(task.cm_data.program_name)


def build_declarations(scenario):
    all_tasks = dict(inspect.getmembers(scenario, inspect.isfunction))
    added = set()
    declarations = []
    for task in all_tasks.values():
        if hasattr(task, IS_CM_TASK):
            _build_declaration(task, all_tasks, added, declarations)
    return declarations


def build_definitions(scenario):
    definitions = []
    for _, task in inspect.getmembers(scenario, inspect.isfunction):
        if hasattr(task, IS_CM_TASK):
            definitions.append(_generate_definition(task))
    return definitions


def main():
    if len(sys.argv) != 2:
        print('Usage: ./update_cm [fast|slow]')
        return
    if sys.argv[1] == 'slow':
        scenario = slow
        filename = '../../clastermaster/datacloud_scenario.sh'
    elif sys.argv[1] == 'fast':
        scenario = fast
        filename = '../../clastermaster/scenario.sh'
    else:
        print('Usage: ./update_cm [fast|slow]')
        return

    definitions = build_definitions(scenario)
    declarations = build_declarations(scenario)

    with open(filename, 'w+') as fh:
        result = template.format(
            full_definition='\n\n'.join(definitions),
            full_declaration='\n\n'.join(declarations),
        )
        fh.write(result)


if __name__ == '__main__':
    main()
