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


class Firmware:
    def __init__(self, vid, pid, cfg, intf, alternate):
        self.mmaps = list()

        self.usbdev = usb.core.find(idVendor = vid, idProduct = pid)
        if self.usbdev is not None:
            self.dfu = dfuse.DfuDevice(self.usbdev)

            mmap_str = usb.util.get_string(self.usbdev, self.dfu.cfg[(0,0)].iInterface)
            mmap = parse.parse('{name}/{address:x}/{regions}',mmap_str)
            self.address = mmap['address']
            for m in mmap['regions'].split(','):
                map = dict(zip(['sectors','sec_size','unit','flag'], parse.parse('{:d}*{:d}{:1}{:1}', m)))
                self.mmaps.append(map)

            for _,alt in self.dfu.alternates():
                if alt.configuration == cfg and alt.bInterfaceNumber == intf and alt.bAlternateSetting == alternate:
                    self.dfu.set_alternate(alt)
                    status = self.dfu.get_status()
                    if status[1] == dfuse.DfuState.DFU_ERROR:
#                        print("Error cleared: %r" % (status,))
                        self.dfu.clear_status() # Clear left-over errors
                    return

        raise ValueError('No DfuSe compatible device found, check device information options (see --help)')

    def erase(self, address):
        print ("Erasing. Please wait this might be long ...")
        """
        dfu.erase(args.erase)
        status = dfu.wait_while_state(dfuse.DfuState.DFU_DOWNLOAD_BUSY)

        if status[1] != dfuse.DfuState.DFU_DOWNLOAD_IDLE:
            raise RuntimeError("An error occured. Device Status: %r" % status)
        """
        for i in range(0, 1280):
            print("Erasing page %r" % i)
            self.dfu.erase(address + i * size)
            status = self.dfu.wait_while_state(dfuse.DfuState.DFU_DOWNLOAD_BUSY)
            if status[1] != dfuse.DfuState.DFU_DOWNLOAD_IDLE:
                raise RuntimeError("An error occured. Device Status: %r" % status)
        print ("Done !")

    def leave_dfu(args):
        dfu = find_device(args)
        dfu.leave()
        status = dfu.get_status()
#    if status[0] > 0:
#        raise RuntimeError("An error occured. Status: %r" % status)

    def flash(dfufile):
        if (dfufile.devInfo['vid'] != dfu.dev.idVendor or dfufile.devInfo['pid'] != self.dfu.dev.idProduct):
            raise ValueError("Vendor/Product id mismatch: [%.4x:%.4x] (file) [%.4x:%.4x] (device). Trying running with --force" % ( \
                dfufile.devInfo['vid'], \
                dfufile.devInfo['vid'], \
                self.dfu.dev.idVendor, \
                self.dfu.dev.idProduct))

        targets = [t for t in dfufile.targets if t['alternate'] == self.dfu.intf.bAlternateSetting]

        if len(targets) == 0:
            raise ValueError("No file target matches the device. Check the --alt setting")


        print ("Flashing. Please wait this might be long ...")
        for t in targets:
            print(self.dfu.alternates()[0][0])
            return

            print ("Found target %r" % t['name'])
            for idx, image in enumerate(t['elements']):
                print("Flashing image %d at 0x%.8X" % (idx, image['address']))

                print("Flashing ...")
                transfer_size = 1024
                self.dfu.set_address(image['address'])
                status = self.dfu.wait_while_state(dfuse.DfuState.DFU_DOWNLOAD_BUSY)
                if status[1] != dfuse.DfuState.DFU_DOWNLOAD_IDLE:
                    raise RuntimeError("An error occured. Device Status: %r" % status)
            
                data = image['data']
                blocks = [data[i:i + transfer_size] for i in range(0, len(data), transfer_size)]
                for blocknum, block in enumerate(blocks):
                    print("Flashing block %r" % blocknum)
                    self.dfu.write(blocknum, block)
                    status = self.dfu.wait_while_state(dfuse.DfuState.DFU_DOWNLOAD_BUSY)
                    if status[1] != dfuse.DfuState.DFU_DOWNLOAD_IDLE:
                        raise RuntimeError("An error occured. Device Status: %r" % status)

                print("Done")

            return

        status = self.dfu.wait_while_state(dfuse.DfuState.DFU_DOWNLOAD_BUSY)

        if status[1] != dfuse.DfuState.DFU_IDLE and status[1] != dfuse.DfuState.DFU_DOWNLOAD_IDLE:
            raise RuntimeError("An error occured. Status: %r" % status)

        print ("Done !")