#!/usr/bin/env python3
"""
ODrive command line utility
"""

from __future__ import print_function
import sys
import argparse
import odrive.discovery
from odrive.utils import Logger, Event

# Flush stdout by default
# Source:
# https://stackoverflow.com/questions/230751/how-to-flush-output-of-python-print
old_print = print
def print(*args, **kwargs):
    kwargs.pop('flush', False)
    old_print(*args, **kwargs)
    file = kwargs.get('file', sys.stdout)
    file.flush() if file is not None else sys.stdout.flush()


## Parse arguments ##
parser = argparse.ArgumentParser(description='ODrive command line utility\n'
                                             'Running this tool without any arguments is equivalent to running `odrivetool shell`\n',
                                 formatter_class=argparse.RawTextHelpFormatter)

# Subcommands
subparsers = parser.add_subparsers(help='sub-command help', dest='command')
shell_parser = subparsers.add_parser('shell', help='Drop into an interactive python shell that lets you interact with the ODrive(s)')
shell_parser.add_argument("--no-ipython", action="store_true",
                          help="Use the regular Python shell "
                          "instead of the IPython shell, "
                          "even if IPython is installed.")

dfu_parser = subparsers.add_parser('dfu', help="Upgrade the ODrive device firmware")
dfu_parser.add_argument('file', metavar='HEX', help='The .hex file to be flashed. Make sure your firmware board version matches the actual board version.')

subparsers.add_parser('liveplotter', help="Upgrade the ODrive's Firmware")
subparsers.add_parser('drv-status', help="Show status of the on-board DRV8301 chips (for debugging only)")
subparsers.add_parser('rate-test', help="Estimate the average transmission bandwidth over USB")
subparsers.add_parser('udev-setup', help="Linux only: Gives users on your system permission to access the ODrive by installing udev rules")

# General arguments
parser.add_argument("-p", "--path", metavar="PATH", action="store",
                    help="The path(s) where ODrive(s) should be discovered.\n"
                    "By default the script will connect to any ODrive on USB.\n\n"
                    "To select a specific USB device:\n"
                    "  --path usb:BUS:DEVICE\n"
                    "usbwhere BUS and DEVICE are the bus and device numbers as shown in `lsusb`.\n\n"
                    "To select a specific serial port:\n"
                    "  --path serial:PATH\n"
                    "where PATH is the path of the serial port. For example \"/dev/ttyUSB0\".\n"
                    "You can use `ls /dev/tty*` to find the correct port.\n\n"
                    "You can combine USB and serial specs by separating them with a comma (no space!)\n"
                    "Example:\n"
                    "  --path usb,serial:/dev/ttyUSB0\n"
                    "means \"discover any USB device or a serial device on /dev/ttyUSB0\"")
parser.add_argument("-s", "--serial-number", action="store",
                    help="The 12-digit serial number of the device. "
                         "This is a string consisting of 12 upper case hexadecimal "
                         "digits as displayed in lsusb. \n"
                         "    example: 385F324D3037\n"
                         "You can list all devices connected to USB by running\n"
                         "(lsusb -d 1209:0d32 -v; lsusb -d 0483:df11 -v) | grep iSerial\n"
                         "If omitted, any device is accepted.")
parser.add_argument("-v", "--verbose", action="store_true",
                    help="print debug information")
parser.add_argument("--version", action="store_true",
                    help="print version information and exit")

parser.set_defaults(path="usb")
args = parser.parse_args()

# Default command
if args.command is None:
    args.command = 'shell'
    args.no_ipython = False

# TODO: deprecate printer - use logger instead
if (args.verbose):
  printer = print
else:
  printer = lambda x: None

logger = Logger(verbose=args.verbose)

def print_version():
    sys.stderr.write("ODrive control utility v" + odrive.__version__ + "\n")
    sys.stderr.flush()

app_shutdown_token = Event()

try:
    if args.version == True:
        print_version()

    elif args.command == 'shell':
        print_version()
        if ".dev" in odrive.__version__:
            print("")
            logger.warn("Developer Preview")
            print("  If you find issues, please report them")
            print("  on https://github.com/madcowswe/ODrive/issues")
            print("  or better yet, submit a pull request to fix it.")
            print("")
        import odrive.shell
        odrive.shell.launch_shell(args, logger, printer, app_shutdown_token)

    elif args.command == 'dfu':
        print_version()
        import odrive.dfu
        odrive.dfu.launch_dfu(args, app_shutdown_token)

    elif args.command == 'liveplotter':
        from odrive.utils import start_liveplotter
        print("Waiting for ODrive...")
        my_odrive = odrive.discovery.find_any(path=args.path, serial_number=args.serial_number)

        # If you want to plot different values, change them here.
        # You can plot any number of values concurrently.
        start_liveplotter(lambda: [my_odrive.motor0.encoder.pos_estimate,
                                my_odrive.motor1.encoder.pos_estimate])

    elif args.command == 'drv-status':
        from odrive.utils import print_drv_regs
        print("Waiting for ODrive...")
        my_odrive = odrive.discovery.find_any(path=args.path, serial_number=args.serial_number)
        print_drv_regs("Motor 0", my_odrive.axis0.motor)
        print_drv_regs("Motor 1", my_odrive.axis1.motor)

    elif args.command == 'rate-test':
        from odrive.utils import rate_test
        print("Waiting for ODrive...")
        my_odrive = odrive.discovery.find_any(path=args.path, serial_number=args.serial_number)
        rate_test(my_odrive)

    elif args.command == 'udev-setup':
        from odrive.utils import setup_udev_rules
        setup_udev_rules(logger)

    else:
        raise Exception("unknown command: " + args.command)

finally:
    app_shutdown_token.set()