#!/usr/bin/python3 -O

# -*- coding: utf-8 -*-

import argparse
from contextlib import redirect_stdout
from inspect import getdoc
from io import StringIO

from colorama import Style

from ereuse_workbench.eraser import EraseType
from ereuse_workbench.tester import Smart
from ereuse_workbench.workbench import Workbench

# Note that this script, if executed directly (./erwb) runs in
# python optimized mode to avoid asserts
if __name__ == '__main__':
    desc = getdoc(Workbench)
    epilog = 'Example: test the computer and save a JSON report to a file: ' \
             'erwb --smart short --stress 2 --json > snapshot.json'
    parser = argparse.ArgumentParser(description=desc, epilog=epilog)
    parser.add_argument('--smart', type=Smart, choices=list(Smart),
                        help='Perform a SMART test to the HDDs.')
    parser.add_argument('--erase', type=EraseType, choices=list(EraseType), help='Erase the HDDs.')
    parser.add_argument('--erase-steps', type=int, default=1, metavar='STEPS',
                        help='Number of erasure STEPS.')
    parser.add_argument('--erase-leading-zeros', action='store_true',
                        help='Shall we perform an extra erasure step writing zeroes?')
    parser.add_argument('--stress', metavar='MINUTES', type=int, default=0,
                        help='Run stress test for the given MINUTES (0 to disable, default)')
    parser.add_argument('--install', type=str,
                        help='The name of the FSA OS to install, without the ".fsa" extension. '
                             'The file has to be in /media/workbench-images')
    parser.add_argument('--server', type=str,
                        help='Connect to a WorkbenchServer at the specified URI. '
                             'This will activate USBSneaky module, load the '
                             'settings from the server and keep indefinitely waiting for an USB to'
                             'be plugged-in.')
    parser.add_argument('--json', action='store_true', help='Print only a JSON report to stdout.')
    args = vars(parser.parse_args())
    print_json = args.pop('json')
    workbench = Workbench(**args)
    if print_json:
        with redirect_stdout(StringIO()):  # We trash all printing from Workbench
            snapshot = workbench.run()
        print(snapshot)
    else:
        workbench.run()
        if args.get('server', None):
            print('You can still link the computer.')
            print('Stop the machine by pressing the power button.')
            print('{}Press CTRL-C to terminate Workbench (you won\'t be able to link).'
                  .format(Style.DIM))
            # We wait indefinitely until the user presses CTRL-C
            # or our child process dies for some reason
            # Note that the child is a daemon so it will be terminated
            # once this main process terminates too
            workbench.usb_sneaky.join()
