import argparse
import importlib.machinery as mach

from rumba.storyboard import *
import rumba.log as log
from rumba.utils import ExperimentManager, PAUSE_SWAPOUT

client1 = Client(
    "rinaperf",
    options="-t perf -s 1000 -c 0",
    c_id='rinaperf_^C'  # To differentiate
)


client2 = Client(
    "rinaperf",
    options="-t perf -s 1000 -D <duration>",
    shutdown="",
    c_id='rinaperf_-D'  # To differentiate
)


server = Server(
    "rinaperf",
    options="-l",
    arrival_rate=0.5,
    mean_duration=5,
    clients=[client1, client2]
)


def main(duration, exp, run=False, script='generated_script.txt'):
    story = StoryBoard(duration)
    story.set_experiment(exp)
    story.add_server_on_node(server, exp.nodes[0])
    client1.add_node(exp.nodes[1])
    if len(exp.nodes) == 2:
        third_n = exp.nodes[1]
    else:
        third_n = exp.nodes[2]
    client2.add_node(third_n)
    story.generate_script()

    with open(script, 'w') as f:
        f.write('################################################\n')
        f.write('# SCRIPT GENERATED WITH RUMBA SCRIPT GENERATOR #\n')
        f.write('################################################\n')
        story.script.write(f)

    if run:
        with ExperimentManager(exp, swap_out_strategy=PAUSE_SWAPOUT):
            exp.swap_in()
            exp.bootstrap_prototype()
            story.start()


if __name__ == '__main__':

    description = "Storyboard generator for Rumba scripts"
    epilog = "2018 Marco Capitani <m.capitani@nextworks.it>"

    parser = argparse.ArgumentParser(description=description,
                                     epilog=epilog)
    parser.add_argument('exp', metavar='EXPERIMENT', type=str,
                        help='The experiment file.')
    parser.add_argument('-O', '--object', metavar='OBJECT', type=str,
                        default='exp',
                        help='The Python variable name of '
                             'the experiment object defined in the file')
    parser.add_argument('d', metavar='DURATION', type=float,
                        help='The required duration of the storyboard')
    parser.add_argument('-R', '--run', action='store_true',
                        help='Run the storyboard after script generation.')
    parser.add_argument('-S', '--script', type=str,
                        default='generated_script.txt',
                        help='Name of the output script.')

    args = parser.parse_args()

    try:
        exp_module = mach.SourceFileLoader('exp', args.exp).load_module()
        experiment = getattr(exp_module, args.object)
        main(args.d, experiment, args.run, args.script)
    except AttributeError:
        logger.error('Module %s has no attribute %s', args.exp, args.object)
    except Exception as e:
        logger.error('Could not load file %s. %s.', args.exp, str(e))
    log.flush_log()