#!/usr/bin/python
#
#  spreadspace pic utils
#
#
#  Copyright (C) 2011 Christian Pointner <equinox@spreadspace.org>
#
#  This file is part of spreadspace pic utils.
#
#  spreadspace pic utils is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  any later version.
#
#  spreadspace pic utils is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with spreadspace pic utils. If not, see <http://www.gnu.org/licenses/>.
#

'''spreadspace simple pic downloader.'''

VERSION_MAJ = 0
VERSION_MIN = 1

### HEX File Magic

def load_hex(file):
    from intelhex import IntelHex

    fin = file
    if fin == '-':
        fin = sys.stdin
    elif not os.path.isfile(fin):
        print >> sys.stderr, "ERROR: File not found: %s" % fin
        sys.exit(1)

    codedata = {}
    ih = IntelHex(fin)
    for a in ih.addresses():
        if a/2 not in codedata.keys():
            codedata[a/2] = 0
        if a%2 == 0:
            codedata[a/2] += ih[a]
        else:
            codedata[a/2] += (ih[a] << 8)

    return codedata

def write_hex(file, codedata):
    from intelhex import IntelHex

    fin = file
    if fin == '-':
        fin = sys.stdout

    hexdata = {}
    for a in codedata:
        hexdata[2*a] = codedata[a] & 0xFF
        hexdata[2*a+1] = codedata[a] >> 8

    ih = IntelHex()
    ih.fromdict(hexdata)
    ih.write_hex_file(fin)


def get_lowest_flash_addr(codedata, fss):
    lowest_code_addr = sorted(codedata.keys())[0]
    return lowest_code_addr - (lowest_code_addr%fss)

def get_highest_flash_addr(codedata, fss):
    highest_code_addr = sorted(codedata.keys())[-1]
    return highest_code_addr + (fss - highest_code_addr%fss)

def create_flash_image(codedata, fss, sa, ea):
    img = ( sa, [0xFFFF]*(ea-sa) )
    for a,d in codedata.items():
        if a < ea:
            img[1][a-sa] = d
    return img

def create_flash_segments(codedata, fs, fss):
    sa = get_lowest_flash_addr(codedata, fss)
    ea = get_highest_flash_addr(codedata, fss)
    if ea >= fs:
        print >> sys.stderr, "WARNING: the hex file contains data after end of flash, these words will be ignored"
        ea = fs

    img = create_flash_image(codedata, fss, sa, ea)
    for i in xrange(0, ea, fss):
        slice = tuple(img[1][i:i+fss])
        if not all( (elem == 0xFFFF) for elem in slice):
            yield (i + img[0], slice)


### Interface to Bootloader

def open_serial(device, baud):
    import os
    import tty
    import termios
    import serial

    print >> sys.stderr,  "opening %s (%s Baud)" % (device, baud)

    try:
        dev = serial.Serial(port=device, baudrate=baud, timeout=3)
        dev.flushInput()
        dev.flushOutput()
        return dev

    except (ValueError, serial.SerialException), msg:
        print >> sys.stderr, "ERROR: opening serial device: %s" % msg
        sys.exit(3)

def calc_csum(str):
    cs = 0
    for c in str:
        cs ^= c
    return cs

def exec_command(dev, cmd, param, answer):
    import struct

    return_codes = { 0: "OK", 1: "invalid command", 2: "bad checksum",
                     3: "not implemented", 4: "flash write error",
                     5: "address invalid", 6: "address prohibited",
                     7: "value out of bounds" }

    dev.flushInput()
    dev.flushOutput()

    cstr = bytearray(struct.pack('<BB', cmd, len(param)+3) + param)
    cstr.extend(struct.pack("<B", calc_csum(cstr)))
    dev.write(cstr)

    astr = bytearray()
    astr += dev.read(4)
    if len(astr) < 4:
        print >> sys.stderr, "ERROR: timeout while reading response header (expected %d bytes, got %d)" % (4, len(astr))
        sys.exit(4)

    if astr[0] != cstr[0]:
        print >> sys.stderr, "ERROR: bootloader returned wrong command code"
        sys.exit(4)

    ret = astr[2]
    if ret != 0:
        rstr = "invalid return code"
        try:
            rstr = return_codes[ret]
        except KeyError:
            pass
        print >> sys.stderr, "ERROR: bootloader returned %d: %s" % (ret, rstr)
        sys.exit(4)

    answer_len = astr[1] - 4
    if answer_len < struct.calcsize(answer):
        print >> sys.stderr, "ERROR: short answer %d bytes received: expected %s bytes" % (answer_len, struct.calcsize(answer))
        sys.exit(4)

    if answer_len > 0:
        tmp = bytearray()
        tmp += dev.read(answer_len)
        if len(tmp) < answer_len:
            print >> sys.stderr, "ERROR: timeout while reading response (expected %d bytes, got %d)" % (answer_len, len(tmp))
            sys.exit(4)

        astr += tmp

    if 0 != calc_csum(astr):
        print >> sys.stderr, "ERROR: checksum error"
        sys.exit(4)

    return struct.unpack_from(answer, astr, 3)

### low level commands

def cmd_identify(dev, name):
    data = exec_command(dev, 1, '', '<BB3sHHBHBH')
    id = { 'ver_min': data[0], 'ver_maj': data[1], 'name': data[2], 'devid': data[3],
           'fs': data[4], 'fss': data[5], 'es': data[6], 'mess': data[7], 'supported': data[8] }

    if id['ver_maj'] != VERSION_MAJ:
        print >> sys.stderr, "incompatible protocol version, expected: %d, got: %d" % (VERSION_MAJ, id['ver_maj'])
        sys.exit(4)
    if name and id['name'] != name:
        print >> sys.stderr, "ERROR: the bootloaders name '%s' differs from the one supplied via" % id['name']
        print >> sys.stderr, "       command line option '%s'. Are sure you are connected to the" % name
        print >> sys.stderr, "       right device?"
        sys.exit(4)

    print >> sys.stderr, "connected with Bootloader '%s' Version %d.%d, (ID=%04X, %d words Flash, FSS=%d, %d bytes EEPROM, MESS=%d)" % \
        (id['name'], id['ver_maj'], id['ver_min'], id['devid'], id['fs'], id['fss'], id['es'], id['mess'])
    return id

def cmd_boot(dev):
    exec_command(dev, 2, '', '<')

def cmd_reset(dev, id):
    exec_command(dev, 3, '', '<')

def cmd_read_flash_segment(dev, id, addr):
    param = struct.pack('<H', addr)
    return exec_command(dev, 4, param, '<%dH' % id['fss'])

def cmd_write_flash_segment(dev, id, addr, data):
    param = struct.pack('<H%dH' % id['fss'], addr, *data)
    return exec_command(dev, 5, param, '<')

def cmd_read_eeprom(dev, id, addr, len):
    param = struct.pack('<HB', addr, len)
    return exec_command(dev, 6, param, '<%dB' % len)

def cmd_write_eeprom(dev, id, addr, data):
    param = struct.pack('<HB%dB' % len(data), addr, len(data), *data)
    return exec_command(dev, 7, param, '<')

def cmd_read_config(dev, id, nr):
    param = struct.pack('<B', nr)
    data = exec_command(dev, 8, param, '<H')
    return data[0]

def cmd_write_config(dev, id, nr, word):
    param = struct.pack('<BH', nr, word)
    return exec_command(dev, 9, param, '<')


### commands
def boot(dev, id, file):
    cmd_boot(dev)

def write_flash(dev, id, file):
    codedata = load_hex(file)
    for segment in create_flash_segments(codedata, id['fs'], id['fss']):
#        print >> sys.stderr, "%05X: %s" % (segment[0], ''.join('%04X '%i for i in segment[1]))
        cmd_write_flash_segment(dev, id, segment[0], segment[1])

def read_flash(dev, id, file):
    codedata = {}
    for addr in xrange(0, id['fs'], id['fss']):
        data = cmd_read_flash_segment(dev, id, addr)
#        print >> sys.stderr, "%05X: %s" % (addr, ''.join('%04X '%i for i in data))
        a = addr
        for d in data:
            codedata[a] = d
            a += 1

    write_hex(file, codedata)

def verify_flash(dev, id, file):
    codedata = load_hex(file)
    err = 0
    for segment in create_flash_segments(codedata, id['fs'], id['fss']):
        flashdata = cmd_read_flash_segment(dev, id, segment[0])
        for file,flash in zip(segment[1] , flashdata):
            if flash == 0x3FFF:
                flash = 0xFFFF
            if flash != file:
                err = 1
                break

        if err !=0:
            break

    if err != 0:
        print >> sys.stderr, "ERROR: verify failed!"
        sys.exit(-1)
    else:
        print >> sys.stderr, "verify ok!"
        sys.exit(0)

commands = {
    'boot': boot,
    'write': write_flash,
    'read': read_flash,
    'verify': verify_flash
}

### Main

if __name__ == '__main__':
    import getopt
    import sys
    import os
    import struct

    usage = '''spreadspace simple pic downloader.
Usage:
    python downloader.py [options] INFILE

Arguments:
    INFILE      name of hex file for downloading.
                Use '-' for reading from stdin.

Options:
    -h, --help              this help message.
    -v, --version           version info.
    --cmd=N                 the command to execute, one out of: "boot,read,write,verify"
    --device=N              the serial port to use (default: /dev/ttyUSB0).
    --baud=N                baudrate to use (default: 57600).
    --name=N                the expected name of the bootloader.
'''

    device = "/dev/ttyUSB0"
    baudrate = 57600
    name = None
    cmd = None

    try:
        opts, args = getopt.getopt(sys.argv[1:], "hv",
                                  ["help", "version", "cmd=", "device=", "baud=", "name="])

        for o, a in opts:
            if o in ("-h", "--help"):
                print >> sys.stderr, usage
                sys.exit(0)
            elif o in ("-v", "--version"):
                print >> sys.stderr, "Version %d.%d" % (VERSION_MAJ, VERSION_MIN)
                sys.exit(0)
            elif o in ("--cmd"):
                cmd = a
            elif o in ("--device"):
                device = a
            elif o in ("--baud"):
                baudrate = a
            elif o in ("--name"):
                name = a

        if not args:
            raise getopt.GetoptError('Hex file is not specified')

        if len(args) > 1:
            raise getopt.GetoptError('Too many arguments')

        if not cmd:
            raise getopt.GetoptError('You have to supply a command (boot,read,write,verify)')

    except getopt.GetoptError, msg:
        print >> sys.stderr, "ERROR: %s" % msg
        print >> sys.stderr, usage
        sys.exit(2)

    dev = open_serial(device, baudrate)
    id = cmd_identify(dev, name)

    try:
        commands[cmd](dev, id, args[0])
    except KeyError:
        print >> sys.stderr, "ERROR: unkown command '%s'" % cmd