import os
from typing import NamedTuple
import argparse

import ihex_loader

class MCU_Config(NamedTuple):
    name: str
    flash_address: int
    flash_size: int
    bootloader_size: int
    signature: int
    page_size: int
    vector_offset: int
    crc_bits: int
    rows_per_block: int

samg55 = MCU_Config('SAMG55', 0x00400000, 0x00080000, 0, 0x1302FD67, 0x200, 0x400, 16, 4)

USB_BOOTLOADER_SIZE = 0x4000 # 16kB for USB bootloader, app starts after this

def load_bin(filename):
    with open(filename, 'rb') as f:
        return f.read()

def generate_lut_entry(lut_index):
    crc_tab_entry = lut_index
    for i in range(8, 0, -1):
        if (crc_tab_entry & 0x01):
            crc_tab_entry = (crc_tab_entry >> 1) ^ 0xEDB88320

        else:
            crc_tab_entry = crc_tab_entry >> 1
    return(crc_tab_entry)

def calc_end_crc(bfr, length):
    crc32 = 0xffffffff

    for i in range(0,length):
        data_byte = bfr[i]

        idx = (data_byte ^ (crc32 & 0x000000FF))

        # work out the CRC
        # the CRC function is the IEEE CRC32 function with a polynomial
        # of 0xEDB88320, implemented using a lookup table
        crc32 = (crc32 >> 8)

        crc32 = crc32 ^ generate_lut_entry(idx)

    return crc32

def check_is_blank_or_ff(input_bytes):
    for bb in input_bytes:
        if(bb != 0 and bb != 0xFF):
            return False

    return True

# Create firmware image for internal ROM bootloader.
# Takes a Hex or Bin file, converts it to a Hex file for 
# direct flashing using a JTAG tool. The new Hex file
# includes the program length vector and CRC at the end
# of the program data.
#
# fwi2cloader_ch341.py can also send the Hex file via I2C
# or SPI bootloader.
def create_fw_image_for_rom_bootloader(input_file, output_file):
    # Check if input_file is recognizable format
    (head,tail) = os.path.split(input_file)
    if(tail.split('.')[-1] == 'hex'):
        # Hex format file
        app_image, startaddr = ihex_loader.load_ihex(input_file)
        if(startaddr != (samg55.flash_address + samg55.bootloader_size)):
            print("Error: Start address doesn't match configuration.")
            return False
        
    elif(tail.split('.')[-1] == 'bin'):
        # Binary format file
        app_image = load_bin(input_file)
        startaddr = samg55.flash_address + samg55.bootloader_size
    else:
        # Can't read any other type
        print("Error: Unknown file type. Only accepts *.bin or *.hex files.")
        return False

    # Fit to 32-byte alignment, but leave 4 bytes at the end for CRC-32
    while (len(app_image) % 32) != 28:
        app_image.append(0xff)

    # Place application length at predetermined slot. Note: this should be empty of any program
    # data. Make sure your linker script skips this section and leaves it empty.
    if(not check_is_blank_or_ff(app_image[samg55.vector_offset:(samg55.vector_offset + 4)])):
        print("Error: Vector location is not blank. Data exists at start location + 0x400")
        return False
    app_image[samg55.vector_offset] = (int(len(app_image)+4) >> 0) & 0xff
    app_image[samg55.vector_offset+1] = (int(len(app_image)+4) >> 8) & 0xff
    app_image[samg55.vector_offset+2] = (int(len(app_image)+4) >> 16) & 0xff
    app_image[samg55.vector_offset+3] = (int(len(app_image)+4) >> 24) & 0xff
    
    # Place calculated CRC in the last 4 bytes, after the length of the application
    crc32 = calc_end_crc(app_image, len(app_image))
    app_image.append(int((crc32 >> 0) & 0xff))
    app_image.append(int((crc32 >> 8) & 0xff))
    app_image.append(int((crc32 >> 16) & 0xff))
    app_image.append(int((crc32 >> 24) & 0xff))

    # Save new Hex file
    ihex_loader.save_ihex(output_file, app_image, startaddr)
    return True

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Creates firmware hex for USB booloader.")
    parser.add_argument('-s', '--suffix', help='suffix to add to filename before the .hex (default=_ROMBL)')
    parser.add_argument('file_name', help='Input filename')
    
    args = parser.parse_args()
    
    (head,tail) = os.path.split(args.file_name)

    if(args.suffix):
        tail_outfile = tail.split('.')[0]+args.suffix+'.hex'
    else:
        tail_outfile = tail.split('.')[0]+'_ROMBL.hex'
    outfile = os.path.join(head,tail_outfile)
    
    passed = create_fw_image_for_rom_bootloader(args.file_name, outfile)

    if(passed):
        print('Rom bootloader file created! {}'.format(outfile))