#!/usr/local/bin/python3
"""
AN1310: A Python bootloader interface for the Microchip AN1310 Bootloader.

Usage:
    from an1310 import AN1310Device
    dev = AN1310Device("/dev/ttyUSB0")
    dev.enter_bootloader()
    dev.read_flash_memory(0, 128)
    dev.run_application()
"""

import argparse
from ctypes import c_ushort
import math
import struct
import serial
import sqlite3
import sys
import time
from PyCRC.CRCCCITT import CRCCCITT
from intelhex import IntelHex

__author__ = 'Stan Janssen'
__url__ = 'https://gitlab.finetuned.nl/bootloader/an1310'
__license__ = 'Public Domain'

__version__ = '1.0'
__status__ = 'Public'

""" Default Values """
BAUDRATE = 115200
PARITY = serial.PARITY_NONE
BYTESIZE = 8
STOPBITS = 1
TIMEOUT = 0.05

STX = 0x0F  # Start of TeXt
ETX = 0x04  # End of TeXt
DLE = 0x05  # Data Link Escape

class AN1310Device():
    def __init__(self, port):
        self.serial = serial.Serial(port=port, baudrate=BAUDRATE,
                                    parity=PARITY, bytesize=BYTESIZE,
                                    stopbits=STOPBITS, timeout=TIMEOUT)
    
    """ Below are the general AN1310 bootloader commands """
    def enter_bootloader(self):
        """ Generate a serial BREAK state to enter bootloader mode """
        print("Sending Break")
        self.serial.read(self.serial.in_waiting)    # flush existing serial buffer
        self.serial.send_break(duration=0.1)        # send break so the PIC goes into bootloader mode
        entered = False
        n = 1
        while entered is False:
            print("Performing autobaud attempt %s" % n)
            n = n + 1
            self.serial.write(bytes([STX]))
            if self.serial.read(1) == b'\x0f':
                entered = True
        print("Successfully entered the bootloader")
        self.read_bootloader_info()
        self._get_device_info()
        
    def read_bootloader_info(self):
        message = [0x00]
        bootloader_info = self._perform_operation(message)
        print(bootloader_info)
        self.bootloader_size = struct.unpack("<H", bytes(bootloader_info[0:2]))[0]
        self.bootloader_version = struct.unpack("<H", bytes(bootloader_info[2:4]))[0]
        
        command_mask_l = bootloader_info[5] & 0x0F      # Get least significant nibble
        if command_mask_l == 2:
            self.pic_type = "PIC16"
        elif command_mask_l == 4:
            self.pic_type = "PIC18"
    
        if self.pic_type == "PIC16":
            command_mask_h = bootloader_info[4]
            if command_mask_h == 0x01:
                self.erase_before_write = False
        else:
            self.erase_before_write = True
        
        self.bootloader_start_address = struct.unpack("<H", bytes(bootloader_info[6:8]))[0]
        if self.pic_type == "PIC16":
            self.deviceid = bootloader_info[10:12]

        return bootloader_info

    def read_flash_memory(self, address, bytes):
        message = [0x01] + self._little_endian(address, "I") + self._little_endian(bytes)
        return self._perform_operation(message)

    def read_flash_memory_crc(self, address, pages):
        message = [0x02] + self._little_endian(address, "I") + self._little_endian(pages)
        return self._perform_operation(message, crc=False)

    def erase_flash_memory(self, page, pages):
        """ Erases a number of pages of flash memory. """
        address = (page + pages) * self.pagesize - 1
        message = [0x03] + self._little_endian(address, "I") + [pages]
        return self._perform_operation(message)

    def write_flash_memory(self, page, pages, data):
        address = page * self.pagesize
        message = [0x04] + self._little_endian(address, "I") + [pages] + data
        return self._perform_operation(message)

    def read_eeprom(self, address, bytes):
        message = [0x05] + self._little_endian(address) + [0x00] + [0x00] + self._little_endian(bytes)
        return self._perform_operation(message)

    def write_eeprom(self, address, bytes, data):
        message = [0x06] + self._little_endian(address) + [0x00] + [0x00] + self._little_endian(bytes) + data
        return self._perform_operation(message)

    def write_config_fuses(self, address, bytes, data):
        message = [0x07] + self._little_endian(address) + [0x00] + [0x00] + [bytes] + data
        return self._perform_operation(message)

    def run_application(self):
        """ Instruct the bootloader to jump to the application firmware """
        message = [0x08]
        self._perform_operation(message, response=False)

    """ Below are some convenience methods for the SmartEVSE """
    def print_cli_menu(self):
        self.serial.write(b'\n')
        time.sleep(0.05)
        print(self.serial.read(self.serial.in_waiting).decode())

    """ Below are methods to verify and program a provided HEX file """
#     def verify_actual_file(self, file):
#         data = IntelHex(file)
#         
#         for segment in data.segments():
#             chip_data = self.read_flash_memory(segment[0], segment[1] - segment[0])
#             
#             file_data = []
#             for i in range(*segment):
#                 file_data.append(data[i])
#                 
#             if file_data != chip_data:
#                 return False
#             else:
#                 print("Segment %s - %s OK" % segment)
#         return True
#     
#     def verify_raw_file(self, file):
#         data = IntelHex(file)
# 
#         """ Read data from the chip """
#         for segment in data.segments():
#             if segment[1] <= self.flashbytes:
#                 range_start = math.floor(segment[0] / self.pagesize)
#                 range_end = math.ceil(segment[1] / self.pagesize)
#                 range_size = range_end - range_start
#                 address_start = range_start * self.pagesize
#                 address_end = range_end * self.pagesize
#             
#                 chip_crcs = self.read_flash_memory_crc(address_start, range_size)
#                 file_crcs = self._file_crc(data, (address_start, address_end))
#                 if chip_crcs != file_crcs:
#                     print("Chunk %s - %s Not OK..." % (address_start, address_end))
#                     print("Chip CRCS were:")
#                     print(chip_crcs)
#                     print("File CRCS were:")
#                     print(file_crcs)
#                     return False
#                 else:
#                     print("Chunk %s - %s OK..." % (address_start, address_end))
#         return True
#         
#     def verify_file(self, file):
#         """ Verifies the file that was programmed using program_file() """
#         data = IntelHex(file)
#         
#         segment = data.segments()[0]
#         first_page = self._page(segment[0])
#         last_page = self._page(min(segment[0], self.bootloader_start_address))
        

    def program_raw_file(self, file):
        """ Programs the contents of the provided hex file straight to the PIC. Does not relocate the reset vectors. May overwrite the bootloader and leave the device in an unusable state. """
        data = IntelHex(file)
        segment = data.segments()[0]
        first_page = self._page(segment[0])
        last_page = self._page(segment[1])
        
        for page in range(first_page, last_page + 1):
            write_data = []
            for j in range(0, self.pagesize):
                write_data.append(data[page * self.pagesize + j])

            new_crc = self._crc(write_data)
            prev_crc = self.read_flash_memory_crc(page * self.pagesize, 1)
            if prev_crc == new_crc:
                print("Nothing to write in page %s, already OK" % page)
            else:
                print("Erasing page %s" % page)
                self.erase_flash_memory(page, 1)
                print("Writing page %s" % page)
                self.write_flash_memory(page, 1, write_data)
                print("Verifying page %s" % page)
                written_crc = self.read_flash_memory_crc(page * self.pagesize, 1)
                if written_crc == new_crc:
                    print("Verified OK")
                else:
                    print("Verified Not OK")

    def program_file(self, file):
        """ Program the provided program, relocating the reset vector in the process. """
        write_size = 16
        
        data = IntelHex(file)
        
        # Get the first two bytes and verify that it is a GOTO statement
        if data[1] != 0xEF:
            raise HEXFileError("No goto-statement in first byte")
        else:
            application_vector = [data[0], data[1], data[2], data[3]]
        
        segment = data.segments()[0]
        first_page = self._page(segment[0])
        last_page = self._page(min(segment[1], self.bootloader_start_address))
        
        print("Programming...")
        st = time.time()
        for page in range(first_page, last_page, write_size):                  # Does not include last_page
            write_data = []
            for j in range(0, self.pagesize * write_size):
                write_data.append(data[page * self.pagesize + j])
            
            if page == 0:
                write_data[0:4] = self._existing_bootloader_pointer()

            existing_crc = self.read_flash_memory_crc(page * self.pagesize, write_size)
            intended_crc = self._chunked_crc(write_data)
            if existing_crc != intended_crc:
                print("Writing pages %s through %s" % (page, page + write_size - 1))
                if self.erase_before_write:
                    self.erase_flash_memory(page, write_size)
                self.write_flash_memory(page, write_size, write_data)
        
        print("Programming complete (%s seconds)" % (time.time() - st))
        
        # Update the application return vector
        app_vector_page = self._page(self.bootloader_start_address) - 1
        print("Relocating the application vector...")
        
        # Update that page:
        app_vector_page_content = self.read_flash_memory(app_vector_page * self.pagesize, self.pagesize)
        if app_vector_page_content[-4:] != application_vector:
            app_vector_page_content[-4:] = application_vector
            print("Erasing page %s", app_vector_page)
            self.erase_flash_memory(app_vector_page, 1)
            print("Writing the following to page %s: %s" %(app_vector_page, app_vector_page_content))
            self.write_flash_memory(app_vector_page, 1, app_vector_page_content)
        else:
            print("Application vector already in place.")

    def _existing_bootloader_pointer(self):
        return self.read_flash_memory(0, 4)     

    def _file_data(self, hexfile, segment):
        data = []
        for i in range(*segment):
            data.append(hexfile[i])
        return data

    def _file_crc(self, hexfile, segment):
        data = self._file_data(hexfile, segment)
        crcs = []
        segment_size = segment[1] - segment[0]
        
        for i in range(0, segment_size, self.pagesize):
            crcs.extend(self._crc(data[i:i+self.pagesize], start=crcs[-2:]))
        return crcs

    """ Below are some internal functions that deal with lower level things """
    def _perform_operation(self, data, response=True, crc=True):
        """ Packages the request in the standard format and sends it to the device"""
        request = bytes([STX] + self._escape_control_chars(data) + self._escaped_crc(data) + [ETX])
        self.serial.write(request)
        if response:
            return self._read_response(crc=crc)

    def _read_response(self, crc=True):
        """ Read the response, un-escaping all control characters, and return the bytestring """
        escape = False
        message = []

        while True:
            if self.serial.in_waiting == 0:
                time.sleep(10/self.serial.baudrate)     # Wait for the next character to be buffered
                continue

            char = ord(self.serial.read(1))             # Read the next char and convert to integer
            if char == DLE and escape is False:
                escape = True
            elif char == ETX and escape is False:       # We've reached the end of the transmission
                message.append(char)
                break
            else:
                escape = False
                message.append(char)

        if self.serial.in_waiting > 0:
            print("Warning: additional bytes left in buffer that were not read:")
            print(self.serial.read(self.serial.in_waiting))

        if crc:
            if self._verify_crc(message) is False:
                print("Warning: CRC check failed")
            return message[1:-3]
        else:
            return message[1:-1]        # The read_flash_crc operation does not include a CRC in the response
    
    def _chunked_crc(self, data):
        crc = []
        
        for d in range(0, len(data), self.pagesize):
            if d == 0:
                crc = self._crc(data[d:d+self.pagesize])
            else:
                crc = crc + self._crc(data[d:d+self.pagesize], start=crc[-2:])
        return crc
    
    def _crc(self, data, start=[0, 0]):
        """ Calculate the CRC and return as two little-endian numbers """
        """ Optionally provide a previous CRC to build on """
        data_as_bytes = bytes(data)
        
        crc_ccitt_tab = [0x0, 0x1021, 0x2042, 0x3063, 0x4084, 0x50a5, 0x60c6, 0x70e7, 0x8108, 0x9129, 0xa14a, 0xb16b, 0xc18c, 0xd1ad, 0xe1ce, 0xf1ef, 0x1231, 0x210, 0x3273, 0x2252, 0x52b5, 0x4294, 0x72f7, 0x62d6, 0x9339, 0x8318, 0xb37b, 0xa35a, 0xd3bd, 0xc39c, 0xf3ff, 0xe3de, 0x2462, 0x3443, 0x420, 0x1401, 0x64e6, 0x74c7, 0x44a4, 0x5485, 0xa56a, 0xb54b, 0x8528, 0x9509, 0xe5ee, 0xf5cf, 0xc5ac, 0xd58d, 0x3653, 0x2672, 0x1611, 0x630, 0x76d7, 0x66f6, 0x5695, 0x46b4, 0xb75b, 0xa77a, 0x9719, 0x8738, 0xf7df, 0xe7fe, 0xd79d, 0xc7bc, 0x48c4, 0x58e5, 0x6886, 0x78a7, 0x840, 0x1861, 0x2802, 0x3823, 0xc9cc, 0xd9ed, 0xe98e, 0xf9af, 0x8948, 0x9969, 0xa90a, 0xb92b, 0x5af5, 0x4ad4, 0x7ab7, 0x6a96, 0x1a71, 0xa50, 0x3a33, 0x2a12, 0xdbfd, 0xcbdc, 0xfbbf, 0xeb9e, 0x9b79, 0x8b58, 0xbb3b, 0xab1a, 0x6ca6, 0x7c87, 0x4ce4, 0x5cc5, 0x2c22, 0x3c03, 0xc60, 0x1c41, 0xedae, 0xfd8f, 0xcdec, 0xddcd, 0xad2a, 0xbd0b, 0x8d68, 0x9d49, 0x7e97, 0x6eb6, 0x5ed5, 0x4ef4, 0x3e13, 0x2e32, 0x1e51, 0xe70, 0xff9f, 0xefbe, 0xdfdd, 0xcffc, 0xbf1b, 0xaf3a, 0x9f59, 0x8f78, 0x9188, 0x81a9, 0xb1ca, 0xa1eb, 0xd10c, 0xc12d, 0xf14e, 0xe16f, 0x1080, 0xa1, 0x30c2, 0x20e3, 0x5004, 0x4025, 0x7046, 0x6067, 0x83b9, 0x9398, 0xa3fb, 0xb3da, 0xc33d, 0xd31c, 0xe37f, 0xf35e, 0x2b1, 0x1290, 0x22f3, 0x32d2, 0x4235, 0x5214, 0x6277, 0x7256, 0xb5ea, 0xa5cb, 0x95a8, 0x8589, 0xf56e, 0xe54f, 0xd52c, 0xc50d, 0x34e2, 0x24c3, 0x14a0, 0x481, 0x7466, 0x6447, 0x5424, 0x4405, 0xa7db, 0xb7fa, 0x8799, 0x97b8, 0xe75f, 0xf77e, 0xc71d, 0xd73c, 0x26d3, 0x36f2, 0x691, 0x16b0, 0x6657, 0x7676, 0x4615, 0x5634, 0xd94c, 0xc96d, 0xf90e, 0xe92f, 0x99c8, 0x89e9, 0xb98a, 0xa9ab, 0x5844, 0x4865, 0x7806, 0x6827, 0x18c0, 0x8e1, 0x3882, 0x28a3, 0xcb7d, 0xdb5c, 0xeb3f, 0xfb1e, 0x8bf9, 0x9bd8, 0xabbb, 0xbb9a, 0x4a75, 0x5a54, 0x6a37, 0x7a16, 0xaf1, 0x1ad0, 0x2ab3, 0x3a92, 0xfd2e, 0xed0f, 0xdd6c, 0xcd4d, 0xbdaa, 0xad8b, 0x9de8, 0x8dc9, 0x7c26, 0x6c07, 0x5c64, 0x4c45, 0x3ca2, 0x2c83, 0x1ce0, 0xcc1, 0xef1f, 0xff3e, 0xcf5d, 0xdf7c, 0xaf9b, 0xbfba, 0x8fd9, 0x9ff8, 0x6e17, 0x7e36, 0x4e55, 0x5e74, 0x2e93, 0x3eb2, 0xed1, 0x1ef0]
        if len(start) == 2:
            crc = struct.unpack("<H", bytes(start))[0]
        else:
            crc = 0

        for d in data:
            tmp = (c_ushort(crc >> 8).value) ^ d
            crc = (c_ushort(crc << 8).value) ^ crc_ccitt_tab[tmp]
        
        return self._little_endian(crc)

    def _escaped_crc(self, data):
        """ Return escaped CRC """
        return self._escape_control_chars(self._crc(data))

    def _verify_crc(self, message):
        """ Verify the received CRC for this message """
        crc = message[-3:-1]
        data = message[1:-3]
        calc_crc = CRCCCITT().calculate(bytes(data))
        if self._little_endian(calc_crc) == crc:
            return True
        else:
            return False

    def _escape_control_chars(self, data):
        """ Escape all control characters in the data, prepending DLE """
        escaped_string = []
        for byte in data:
            if byte in [STX, ETX, DLE]:
                escaped_string = escaped_string + [DLE]
            escaped_string = escaped_string + [byte]
        return escaped_string

    def _little_endian(self, value, format="H"):
        """ Return a little-endian list representation of the provided value """
        return list(struct.pack("<" + format, value))

    def _load_eeprom_data(self, hexfile):
        pass
        
    def _page(self, address):
        return int(address / self.pagesize)
    
    def _get_device_info(self):
        if self.pic_type == "PIC18":
            deviceid = struct.unpack("<H", bytes(self.read_flash_memory(0x3FFFFE, 2)))[0] >> 5
        print("The device ID is %s" % deviceid)
        devicedb = sqlite3.connect("devices.db")
        cur = devicedb.cursor()
        cur.execute("SELECT `ENDFLASH` - `STARTFLASH`, `WRITEFLASHBLOCKSIZE`, `ERASEFLASHBLOCKSIZE`, `STARTEE`, `ENDEE` - `STARTEE`, `ENDGPR` - `STARTGPR` FROM `devices` WHERE `DEVICEID` = ?", (deviceid,))
        result = cur.fetchall()
        if len(result) > 0:
            self.flashbytes, self.pagesize, self.erase_pagesize, self.eeprom_start_address, self.eeprom_size, self.ram_size = result[0]
        else:
            raise Exception("Device ID %s was not found in devices.db" % deviceid)
        cur.close()

class HEXFileError(Exception):
    def __init__(self, message):
        self.message = message
    
    def __str__(self):
        return repr(self.message)
    

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="AN1310 Bootloader application for PIC Microchip controllers.")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--program", "-p", help="Program a file", action="store_true")
    group.add_argument("--verify", "-v", help="Verify a file", action="store_true")
    group.add_argument("--program-raw", "-r", help="Program a raw, compiled file", action="store_true")
    parser.add_argument("--device", "-d", nargs=1, help="Device (serial port)", required=True)
    parser.add_argument("--baudrate", "-b", nargs=1, type=int, help="baudrate")
    parser.add_argument("FILE", nargs=1, help="File input")
    options = parser.parse_args()
    if options.program:
        dev = AN1310Device(options.device[0])
        dev.enter_bootloader()
        dev.program_file(options.FILE[0])
        dev.run_application()
    elif options.verify:
        dev = AN1310Device(options.device[0])
        dev.enter_bootloader()
        dev.verify_file(options.FILE[0])
        dev.run_application()
    elif options.program_raw:
        dev = AN1310evice(options.device[0])
        dev.enter_bootloader()
        dev.program_raw_file(options.FILE[0])
        dev.run_application()