#!/usr/bin/env python3
import argparse
import usb.core
import usb.util
import sys
import time

def load_bootimg_to_device(binfile, dev):
    if dev.is_kernel_driver_active(0):
        try:
            dev.detach_kernel_driver(0)
        except usb.core.USBError as e:
            print("\tCould not detach kernel driver from interface({0}): {1}".format(0, str(e)))
            return
    #dev.reset()
    dev.set_configuration()

    cfg = dev[0]
    intf = cfg[(0,0)]
    in_ep = intf[0]
    out_ep = intf[1]

    print(out_ep)

    fwbin = open(binfile, "rb").read()
    sz = len(fwbin)
    fwbin = bytearray(fwbin)
    print('\n\tFilename: ', binfile, ", Size: ", sz)

    # Byteswap 16bit words
    fwbin[0::2], fwbin[1::2] = fwbin[1::2], fwbin[0::2]
    for ofs in range(0, len(fwbin), 64):
        out_ep.write(fwbin[ofs:ofs+64])

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
                    prog='ASI3500_usbboot',
                    description='ASI3500 USB boot utility')
    parser.add_argument('binfile', default='bootimg.bin')
    parser.add_argument('-i', '--index', default=None)
    parser.add_argument('-l', '--loop', action='store_true')
    args = parser.parse_args()

    binfile = args.binfile
    selected_index = args.index
    loop = args.loop

    spinner = '|/-\\'
    spinner_pos = 0

    print("Looking for ASI3500 devices:")
    while True:
        if loop:
            print('\033[1D', end='')
            print(spinner[spinner_pos % len(spinner)], end='', flush=True)
        spinner_pos += 1
        # Find our device(s)
        time.sleep(1.0)
        devices = list(usb.core.find(find_all=True, idVendor=0x0451, idProduct=0x9010))
        if not len(devices):
            continue
        print('\033[1D', end='')
        print("Found {} matching devices".format(len(devices)))
        for index, dev in enumerate(devices):
            if selected_index == None or selected_index == index:
                print("loading {} to index {}".format(binfile, index))
                load_bootimg_to_device(binfile, dev)
            else:
                print("skipping index {}".format(index))
        if not loop:
            break
