import serial
import sys
import time
import binascii
import random
import pprint
import collections
import struct

import elffile

class Error(Exception):
    pass

RETURN_CODE = {
    0: 'CMD_SUCCESS',
    1: 'INVALID_COMMAND',
    2: 'SRC_ADDR_ERROR',
    3: 'DST_ADDR_ERROR',
    4: 'SRC_ADDR_NOT_MAPPED',
    5: 'DST_ADDR_NOT_MAPPED',
    6: 'COUNT_ERROR',
    7: 'INVALID_SECTOR',
    8: 'SECTOR_NOT_BLANK',
    9: 'SECTOR_NOT_PREPARED_FOR_WRITE_OPERATION',
    10: 'COMPARE_ERROR',
    11: 'BUSY',
    12: 'PARAM_ERROR',
    13: 'ADDR_ERROR',
    14: 'ADDR_NOT_MAPPED',
    15: 'CMD_LOCKED',
    16: 'INVALID_CODE',
    17: 'INVALID_BAUD_RATE',
    18: 'INVALID_STOP_BIT',
    19: 'CODE_READ_PROTECTION_ENABLED ',
}

Section = collections.namedtuple('Section', 'name address content')

def fix_vectors(data):
    """Patch the vector table so the bootloader recognises the code
    as valid.  Sum of the first eight is zero.
    """
    first = data[:4*7]
    remainder = data[4*8:]

    vectors = struct.unpack('<IIIIIII', first)
    last = -sum(vectors) & 0xFFFFFFFF

    return first + struct.pack('<I', last) + remainder

class ELFReader:
    """Reads an ELF file into sections."""
    def __init__(self, name):
        elf = elffile.open(name)

        self.entry = elf.fileHeader.entry
        self.sections = []

        for header in elf.sectionHeaders:
            if header.name in ['.text', '.data', '.vectors']:
                content = header.content

                if header.name == '.vectors':
                    content = fix_vectors(content)

                self.sections.append(Section(header.name, header.addr, content))
        
class Loader:
    """Talk with a LPC1114 over the serial bootloader protocol."""

    SYNCHRONIZED = 'Synchronized\r\n'
    OK = 'OK\r\n'

    def __init__(self, port='/dev/ttyUSB0', verbose=True):
        self.verbose = verbose

        self.port = serial.Serial(port, 115200, timeout=0.5)

    def log(self, msg, eol=True):
        if self.verbose:
            if eol:
                print msg
            else:
                print msg,

    def set_reset(self, level):
        self.port.setDTR(level == 0)

    def set_bsl(self, level):
        self.port.setRTS(level == 0)

    def run(self):
        self.set_bsl(1)
        self.set_reset(1)
        self.set_reset(0)
        self.set_reset(1)

    def sync(self):
        for tries in range(3):
            self.set_bsl(0)
            self.set_reset(1)
            self.set_reset(0)
            self.set_reset(1)

            self.read(100)
        
            self.write('?')

            got = self.read(len(self.SYNCHRONIZED))

            if got == self.SYNCHRONIZED:
                self.send('Synchronized')
                got = self.read(len(self.OK))

                if got == self.OK:
                    self.send('12000')

                    got = self.read(len(self.OK))
                    if got == self.OK:
                        self.log(self.execute('N', (int, int, int, int)))
                        return

    def load(self, address, data):
        chunk = 512
        per_line = 45

        for offset in range(0, len(data), 512):
            part = data[offset:offset+512]

            self.send('W %d %d' % (address + offset, len(part)))
            self.check_ok()

            for micro in range(0, len(part), per_line):
                encoded = binascii.b2a_uu(part[micro:micro+per_line])
                assert encoded.endswith('\n')
                encoded = encoded[:-1]

                self.send(encoded)

            check = sum(ord(x) for x in part)
            self.send('%d' % check)

            got = self.readline()

            if got == 'OK':
                pass
            elif got == 'RESEND':
                raise Error('Checksum error on write')
            else:
                raise Error('Protocol error')

    def prepare(self, first, last=None):
        if last == None:
            last = first

        self.execute('P', (), first, last)

    def erase(self, first, last=None):
        if last == None:
            last = first

        self.prepare(first, last)
        self.execute('E', (), first, last)

    def program(self, address, data):
        page_size = 256
        chunk = 512
        load_address = 0x10000400

        if address % page_size != 0:
            raise Error('Address must be page size aligned')

        remainder = len(data) % page_size

        if remainder != 0:
            data += '\xff' * (page_size - remainder)

        assert len(data) % page_size == 0

        for offset in range(0, len(data), chunk):
            part = data[offset:offset+chunk]
            self.load(load_address, part)
            self.prepare(0)
            self.execute('C', (), address + offset, load_address, len(part))

    def go(self, address):
        self.execute('G', (), address, 'T')

        for tries in range(5):
            try:
                self.execute('J', (int, ))
                break
            except Error:
                self.read(100)
                time.sleep(0.1)

    def unlock(self):
        self.execute('U', (), 23130)

    def check_ok(self):
        code = self.readline()

        try:
            code = int(code)
        except ValueError:
            raise Error('Invalid result code')

        if code != 0:
            raise Error(RETURN_CODE[code])

    def execute(self, command, result, *kargs):
        send = command

        for arg in kargs:
            send += ' %s' % arg

        self.send(send)
        self.check_ok()

        values = []

        for t in result:
            values.append(t(self.readline()))

        return values

    def readline(self):
        v = ''

        self.log('<', eol=False)

        try:
            while True:
                got = self.read(1, True)

                if got == '':
                    raise Error('No response')

                if got == '\r':
                    got = self.read(1, True)

                    if got == '\n':
                        return v
                else:
                    v += got
        finally:
            self.log('')

    def write(self, v):
        self.log('> %r' % v)
        self.port.write(v)

    def send(self, v):
        v += '\r\n'

        self.write(v)
        got = self.read(len(v))

        if got != v:
            raise Error('Echo error.  Sent %r, got %r' % (v, got))

    def read(self, limit, partial=False):
        got = self.port.read(limit)

        if partial:
            self.log('%r' % got, eol=False)
        else:
            self.log('< %r' % got)

        return got

def main():
    loader = Loader()
    loader.sync()
    loader.unlock()

    loader.erase(0)

    for arg in sys.argv[1:]:
        data = ELFReader(arg)

        for section in data.sections:
            loader.program(section.address, section.content)

    loader.sync()
    loader.run()
#    loader.unlock()
#    loader.go(data.entry)

if __name__ == '__main__':
    main()