#!D:\Program Files\Python312\python.exe
"""
ODrive command line utility
"""

from __future__ import print_function
import sys

if sys.version_info < (3, 6):
    print("Your Python version (Python {}.{}) is too old. Please install Python 3.6 or newer.".format(
        sys.version_info.major, sys.version_info.minor
    ))
    exit(1)

import asyncio
import sys
import os
import argparse
import tempfile
import time
import math
import platform

import mwdrive
from mwdrive.utils import OperationAbortedException
from fibre import Logger

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

script_path=os.path.dirname(os.path.realpath(__file__))

## 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.")

# These two DFU commands are aliases except for the info message
for command in ['dfu', 'legacy-dfu']:
    dfu_parser = subparsers.add_parser(command, help="Upgrade the ODrive device firmware."
                                                "If no serial number is specified, the first ODrive that is found is updated")
    dfu_parser.add_argument('file', nargs='?',
                            help='The .elf file to be flashed. Make sure target board version '
                            'of the firmware file matches the actual board version. '
                            'You can download the latest release manually from '
                            'https://github.com/odriverobotics/ODrive/releases. '
                            'If no file is provided, the script automatically downloads '
                            'the firmware from the internet.')
    dfu_parser.add_argument('--channel', action='store',
                            help='The release channel to download from. '
                            'Mutually exclusive with --version. '
                            'If no channel, no version and no firmware file is specified, '
                            'the "master" channel is used. ')
    dfu_parser.add_argument('--version', action='store',
                            help='The version name to download. '
                            'Mutually exclusive with --channel.')
    dfu_parser.add_argument('--no-erase-all', action='store_true',
                            help='If specified, only those flash sectors occupied '
                            'by the firmware will be erased. If not specified, all '
                            'flash sectors (including non-volatile config) will be '
                            'erased.')

# Vendor DFU firmware upgrade (BIN) using vendor tools (Windows)
upgrade_parser = subparsers.add_parser(
    'upgrade',
    help="Upgrade MWDrive firmware using vendor DFU tools (HPM / Nations GM).",
)
upgrade_parser.add_argument(
    'bin',
    help="Path to the .bin firmware image to flash.",
)
upgrade_parser.add_argument(
    '--vendor',
    choices=['auto', 'hpm', 'gm'],
    default='auto',
    help="Select vendor toolchain. Default: auto (detect HPM/GM DFU device).",
)
upgrade_parser.add_argument(
    '--erase-all',
    action='store_true',
    help="Erase full flash before flashing (if supported by the vendor tool).",
)
upgrade_parser.add_argument(
    '--assets-dir',
    default=None,
    help="Path to assets dir that contains vendor tools (e.g. odrive-firmware-upgrader/assets).",
)
upgrade_parser.add_argument(
    '--hpm-dfu-exe',
    default=None,
    help="Path to hpm_dfu.exe (overrides --assets-dir and env).",
)
upgrade_parser.add_argument(
    '--gm-tool-dir',
    default=None,
    help="Path to GM tool directory containing NsMcuCmdDll.dll (overrides --assets-dir and env).",
)
upgrade_parser.add_argument(
    '--hpm-usb-id',
    default="0x34b7,0x0006",
    help="HPM DFU USB ID used by hpm_dfu.exe, format '0xVID,0xPID'.",
)

# Erase-all only (no firmware flash, no reset)
erase_all_parser = subparsers.add_parser(
    'erase-all',
    help="Enter vendor DFU mode and erase full flash without programming or reset.",
)
erase_all_parser.add_argument(
    '--vendor',
    choices=['auto', 'hpm', 'gm'],
    default='auto',
    help="Select vendor toolchain. Default: auto (detect HPM/GM DFU device).",
)
erase_all_parser.add_argument(
    '--assets-dir',
    default=None,
    help="Path to assets dir that contains vendor tools (e.g. odrive-firmware-upgrader/assets).",
)
erase_all_parser.add_argument(
    '--hpm-dfu-exe',
    default=None,
    help="Path to hpm_dfu.exe (overrides --assets-dir and env).",
)
erase_all_parser.add_argument(
    '--gm-tool-dir',
    default=None,
    help="Path to GM tool directory containing NsMcuCmdDll.dll (overrides --assets-dir and env).",
)
erase_all_parser.add_argument(
    '--hpm-usb-id',
    default="0x34b7,0x0006",
    help="HPM DFU USB ID used by hpm_dfu.exe, format '0xVID,0xPID'.",
)

unlock_parser = subparsers.add_parser('unlock', help="Try to remove read-out protection."
                                               "If no serial number is specified, the first ODrive that is found is unlocked")

backup_config_parser = subparsers.add_parser('backup-config', help="Saves the configuration of the ODrive to a JSON file")
backup_config_parser.add_argument('file', nargs='?',
                        help="Path to the file where to store the data. "
                        "If no path is provided, the configuration is stored in {}.".format(tempfile.gettempdir()))

restore_config_parser = subparsers.add_parser('restore-config', help="Restores the configuration of the ODrive from a JSON file")
restore_config_parser.add_argument('file', nargs='?',
                        help="Path to the file that contains the configuration data. "
                        "If no path is provided, the configuration is loaded from {}.".format(tempfile.gettempdir()))

subparsers.add_parser('liveplotter', help="For plotting of odrive parameters (i.e. position) in real time")
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")

# 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:idVendor=0x1209,idProduct=0x0D32,bInterfaceClass=0,bInterfaceSubClass=1,bInterfaceProtocol=0")
parser.set_defaults(path="usb:idVendor=0x1209,idProduct=0x0D32")
args = parser.parse_args()

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

logger = Logger(verbose=args.verbose)

def print_version():
    sys.stderr.write("CyberBeast Motor Driver control utility v" + mwdrive.__version__ + "\n")
    sys.stderr.flush()

if platform.system() == 'Linux':
    if not os.path.isfile('/etc/udev/rules.d/91-mwdrive.rules'):
        logger.warn(
            "Device permissions are not set up. Please exit odrivetool and run the following command: \n"
            "sudo bash -c \"curl https://cdn.odriverobotics.com/files/odrive-udev-rules.rules > /etc/udev/rules.d/91-mwdrive.rules && udevadm control --reload-rules && udevadm trigger\"\n"
            "otherwise your ODrive may not be recognized.")

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

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

    elif args.command == 'dfu' or args.command == 'legacy-dfu':
        print_version()

        if args.command == 'dfu':
            logger.warn("Try our new firmware update system!")
            logger.warn("Supports firmware updates from the Web GUI and firmware updates via CAN bus.")
            logger.warn("More info and migration instructions here: https://docs.odriverobotics.com/v/devel/guides/new-dfu.html.")

        import mwdrive.dfu
        try:
            if 'win32' in sys.platform:
                # needed for subprocess support
                # https://docs.python.org/3/library/asyncio-platforms.html#asyncio-windows-subprocess
                asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())

            if sum([bool(args.file), bool(args.channel), bool(args.version)]) > 1:
                raise mwdrive.dfu.DfuError("Only one of firmware-file, --channel or --version must be specified.")
            elif sum([bool(args.file), bool(args.channel), bool(args.version)]) == 0:
                args.channel = 'master'

            asyncio.run(mwdrive.dfu.launch_dfu(
                serial_number=args.serial_number,
                path=args.file,
                channel=args.channel,
                version=args.version,
                erase_all=not args.no_erase_all,
                logger=logger))
        except mwdrive.dfu.DfuError as ex:
            logger.error(str(ex))
            sys.exit(1)

    elif args.command == 'upgrade':
        print_version()
        if platform.system() != 'Windows':
            logger.error("The 'upgrade' command is currently supported on Windows only.")
            sys.exit(1)
        from mwdrive.firmware_upgrade import upgrade_firmware
        try:
            upgrade_firmware(
                bin_path=args.bin,
                vendor=args.vendor,
                erase_all=args.erase_all,
                path=args.path,
                serial_number=args.serial_number,
                assets_dir=args.assets_dir,
                hpm_dfu_exe=args.hpm_dfu_exe,
                gm_tool_dir=args.gm_tool_dir,
                hpm_usb_id=args.hpm_usb_id,
            )
        except Exception as ex:
            logger.error(str(ex))
            sys.exit(1)

    elif args.command == 'erase-all':
        print_version()
        if platform.system() != 'Windows':
            logger.error("The 'erase-all' command is currently supported on Windows only.")
            sys.exit(1)
        from mwdrive.firmware_upgrade import erase_all_flash_only
        try:
            erase_all_flash_only(
                vendor=args.vendor,
                path=args.path,
                serial_number=args.serial_number,
                assets_dir=args.assets_dir,
                hpm_dfu_exe=args.hpm_dfu_exe,
                gm_tool_dir=args.gm_tool_dir,
                hpm_usb_id=args.hpm_usb_id,
            )
        except Exception as ex:
            logger.error(str(ex))
            sys.exit(1)

    elif args.command == 'unlock':
        print_version()
        import mwdrive.dfu
        mwdrive.dfu.unlock_device(args.serial_number, None)

    elif args.command == 'liveplotter':
        from mwdrive.utils import start_liveplotter
        print("Waiting for ODrive...")
        my_odrive = mwdrive.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.
        cancellation_token = start_liveplotter(lambda: [
            my_mdrive.axis0.encoder.pos_estimate,
            my_mdrive.axis1.encoder.pos_estimate,
        ])

        print("Showing plot. Press Ctrl+C to exit.")
        while not cancellation_token.is_set():
            time.sleep(1)

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

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

    elif args.command == 'backup-config':
        from mwdrive.legacy_config import backup_config
        print("Waiting for ODrive...")
        my_odrive = mwdrive.find_any(path=args.path, serial_number=args.serial_number)
        backup_config(my_odrive, args.file, logger)

    elif args.command == 'restore-config':
        from mwdrive.legacy_config import restore_config
        print("Waiting for ODrive...")
        my_odrive = mwdrive.find_any(path=args.path, serial_number=args.serial_number)
        restore_config(my_odrive, args.file, logger)

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

except OperationAbortedException:
    logger.info("Operation aborted.")
