from rumboot.ops.base import base
from rumboot.ops.xfer import basic_uploader
import tqdm
import rumboot.xfer


class bootdownload(object):
    '''
    Hisilicon boot downloader

    >>> downloader = bootdownload()
    >>> downloader.burn(filename)

    '''
    def init_bootmode(self):
        i = 0
        counter = 0
        while i < 30:
            in_bin = self.s.read(1)
            if in_bin == b'\x00':
                continue
            if in_bin == b'\x20':
                counter = counter + 1
            if counter == 5:
                self.s.flushOutput()
                self.s.write(settings.AA)
                self.s.flushInput()
                print("Welcome to boot-mode\n")
                return None
            i = i + 1
        sys.exit(1)

    def __init__(self, chiptype, serialport='/dev/ttyUSB0'):
        print("Trying open", serialport)
        try:
            self.s = serial.Serial(
                port=serialport,
                baudrate=115200,
                parity=serial.PARITY_NONE,
                stopbits=serial.STOPBITS_ONE,
                bytesize=serial.EIGHTBITS,
                timeout=None
                #timeout = 1
            )
            self.init_bootmode()

        except serial.serialutil.SerialException:
            #no serial connection
            self.s = None
            print("\nFailed to open serial!", serialport)
            sys.exit(2)

        with open(chiptype+".json", "r") as read_file:
            data = json.load(read_file)
        print(data)
        self.chip = data

    def __del__(self):
        print("Exiting...")
        if self.s is not None:
            self.s.close()

    def calc_crc(self, data, crc=0):
        for char in data:
            crc = ((crc << 8) | char) ^ settings.crctable[(crc >> 8) & 0xff]
        for i in range(0, 2):
            crc = ((crc << 8) | 0) ^ settings.crctable[(crc >> 8) & 0xff]
        return crc & 0xffff

    def sendframe(self, data, loop):
        for i in range(1, loop):
            self.s.flushOutput()
            if self.DEBUG == 1:
                if len(data) > 20:
                    print(
                        "len: ",
                        len(data),
                        "write : [",
                        ''.join('%02x ' % i for i in data[:20]), "... ]"
                    )
                else:
                    print(
                        "len: ",
                        len(data),
                        "write : [",
                        ''.join('%02x ' % i for i in data), "]"
                    )
            self.s.write(data)
            self.s.flushInput()
            try:
                ack = self.s.read()
                if len(ack) == 1:
                    if self.DEBUG == 1:
                        print(
                            "ret ack   : ",
                            ''.join('0x%02x ' % i for i in ack)
                        )
                    if ack == settings.ack_b:
                        return None
            except:
                return None

        print('failed')

    def send_head_frame(self, length, address):
        print('Send HEAD frame...')
        self.s.timeout = 0.03

        frame = bytearray(14)
        frame[0] = settings.twos_complement_to_int(-2, 8)
        frame[1] = settings.twos_complement_to_int(0, 8)
        frame[2] = settings.twos_complement_to_int(-1, 8)
        frame[3] = settings.twos_complement_to_int(1, 8)

        frame[4] = (length >> 24) & 0xff
        frame[5] = (length >> 16) & 0xff
        frame[6] = (length >> 8) & 0xff
        frame[7] = (length) & 0xff
        frame[8] = (address >> 24) & 0xff
        frame[9] = (address >> 16) & 0xff
        frame[10] = (address >> 8) & 0xff
        frame[11] = (address) & 0xff
        crc = self.calc_crc(frame[:12])
        frame[12] = (crc >> 8) & 0xff
        frame[13] = crc & 0xff

        self.sendframe(frame, 16)

    def send_ddrstep(self):
        print('Send DDRSTEP frame...')
        seq = 1
        self.s.timeout = 0.15

        address0 = int(self.chip["ADDRESS"][0], 16)
        self.send_head_frame(64, address0)
        head = bytearray(3)
        head[0] = 0xDA  # b'\xDA'  # 0xDA
        head[1] = seq & 0xFF
        head[2] = (~seq) & 0xFF
        data = head + bytes(self.chip["DDRSTEP0"])
        crc = self.calc_crc(data)
        h = bytearray(2)
        h[0] = (crc >> 8) & 0xff
        h[1] = crc & 0xff
        data += h

        self.sendframe(data, 16)
        self.send_tail_frame(seq + 1)

    def send_tail_frame(self, seq):
        print('Send TAIL frame...')
        data = bytearray(3)
        data[0] = 0xED
        data[1] = seq & 0xFF
        data[2] = (~seq) & 0xFF

        crc = self.calc_crc(data)
        h = bytearray(2)
        h[0] = (crc >> 8) & 0xff
        h[1] = crc & 0xff
        data += h

        self.sendframe(data, 16)

    def send_data_frame(self, seq, data):
        self.s.timeout = 0.15
        head = bytearray(3)
        head[0] = 0xDA
        head[1] = seq & 0xFF
        head[2] = (~seq) & 0xFF
        data = head + data

        crc = self.calc_crc(data)
        h = bytearray(2)
        h[0] = (crc >> 8) & 0xff
        h[1] = crc & 0xff
        data += h

        self.sendframe(data, 32)

    def store_SPL(self, data):
        first_length = int(self.chip["FILELEN"][1], 16)
        address1 = int(self.chip["ADDRESS"][1], 16)
        self.send_head_frame(first_length, address1)

        bar = progressbar.ProgressBar(maxval=first_length, widgets=[
            'Send DATA frame',
            progressbar.Bar(left='[', marker='=', right=']'),
            progressbar.SimpleProgress(),
        ]).start() if self.DEBUG == 0 else None

        seq = 1
        data = data[:first_length]
        for chunk in (
            data[_:_+settings.MAX_DATA_LEN] for _ in range(
                0, first_length, settings.MAX_DATA_LEN
            )
        ):
            self.send_data_frame(
                seq,
                chunk
            )
            bar.update(seq*len(chunk)) if self.DEBUG == 0 else None
            seq += 1
        bar.finish() if self.DEBUG == 0 else None
        self.send_tail_frame(seq)

    def store_Uboot(self, data):
        length = len(data)
        address2 = int(self.chip["ADDRESS"][2], 16)
        self.send_head_frame(length, address2)

        bar = progressbar.ProgressBar(maxval=length, widgets=[
            'Send DATA frame',
            progressbar.Bar(left='[', marker='=', right=']'),
            progressbar.SimpleProgress(),
        ]).start() if self.DEBUG == 0 else None
        seq = 1
        for chunk in (
            data[_:_+settings.MAX_DATA_LEN] for _ in range(
                0, length, settings.MAX_DATA_LEN
            )
        ):
            self.send_data_frame(
                seq,
                chunk
            )
            bar.update(seq*len(chunk)) if self.DEBUG == 0 else None
            seq += 1
        bar.finish() if self.DEBUG == 0 else None
        self.send_tail_frame(seq)

    def send_data(self, data):
        self.send_ddrstep()
        self.store_SPL(data)
        self.store_Uboot(data)

    def burn(self, filename):
        f = open(filename, "rb")
        data = f.read()
        f.close()

        print('Sending', filename, '...')
        self.send_data(data)
        print('Done\n')

class hisi_chips_uploader(basic_uploader):
    def on_start(self):
        #HACK: Old chips don't print anything via uart
        #HACK: edcl is the only way to bring 'em up
        #HACK: We do the lucky guess by looking at the image header
        def prg(total_bytes, position, increment):
                self.term.progress.update(increment)     

        tp = self.term.formats.guess(self.term.runlist[0])
        if self.term.hack("silentRom"):
            print("WARNING: Bootloader doesn't support xmodem, forcing initial edcl upload")
            old = self.term.xfer.how
            self.term.xfer.selectTransport("edcl")
            if not self.term.xfer.push(self.term.chip.spl_address):
                return 1
            self.term.xfer.selectTransport(old)

        if not tp:
            return True

        if tp.name == "Legacy K1879XB1YA":
            self.term.xfer.write32(0x00100010, 0x00100014)
            return True
        if tp.name == "NM6408 (Legacy)":
            self.term.xfer.write32(0x30000, 0x1)
            return True