#!/usr/bin/env python3

import argparse
import sys
import pathlib
import zlib

missing_modules = []

try:
    from bitstring import BitArray
except ImportError:
    missing_modules.append("bitstring")

try:
    import construct
    from construct import this, len_, Struct, Padded, CString, Rebuild, Computed, RawCopy, Checksum, Hex, Bytes, Int8ul, Int16ul, Int32ul, Const, Flag, Padding, Array, Prefixed, GreedyBytes, GreedyRange, PaddedString
except ImportError:
    missing_modules.append("construct")

if len(missing_modules) > 0:
    error_text = ""
    if 'build' in sys.argv:
        error_text += """
Warning: no dfu file has been generated (missing python3 dependencies)
"""
    error_text += f"""
dfu requires the following python module(s): {', '.join(missing_modules)}
install with: python3 -m pip install {' '.join(missing_modules)}
"""
    # Soft fail to avoid cmake error
    print(error_text)
    sys.exit(0)


DFU_SIGNATURE = b'DfuSe'
DFU_size = Rebuild(Int32ul, 0)


def DFU_file_length(ctx):
    '''Compute the entire file size + 4 bytes for CRC

    The total DFU file length is ostensibly the actual
    length in bytes of the resulting file.

    However DFU File Manager does not seem to agree,
    since it's output size is 16 bytes short.

    Since this is suspiciously the same as the suffix
    length in bytes, we omit that number to match
    DFU File Manager's output.

    '''
    size = 11        # DFU Header Length
    # size += 16     # DFU Suffix Length
    for target in ctx.targets:
        # Each target has a 274 byte header consisting
        # of the following fields:
        size += Const(DFU_SIGNATURE).sizeof()          # szSignature ('Target' in bytes)
        size += Int8ul.sizeof()                        # bAlternateSetting
        size += Int8ul.sizeof()                        # bTargetNamed
        size += Padding(3).sizeof()                    # Padding
        size += Padded(255, CString('utf8')).sizeof()  # szTargetName
        size += Int32ul.sizeof()                       # dwTargetSize
        size += Int32ul.sizeof()                       # dwNbElements
        size += DFU_target_size(target)

    return size


def DFU_target_size(ctx):
    '''Returns the size of the target binary data, plus the
    dwElementAddress header, and dwElementSize byte count.
    '''
    size = 0

    try:
        images = ctx.images
    except AttributeError:
        images = ctx['images']

    size += sum([DFU_image_size(image) for image in images])
    return size


def DFU_image_size(image):
    return len(image['data']) + Int32ul.sizeof() + Int32ul.sizeof()


DFU_image = Struct(
    'dwElementAddress' / Hex(Int32ul),     # Data offset address for image
    'data' / Prefixed(Int32ul, GreedyBytes)
)

DFU_target = Struct(
    'szSignature' / Const(b'Target'),      # DFU target identifier
    'bAlternateSetting' / Int8ul,          # Gives device alternate setting for which this image can be used
    'bTargetNamed' / Flag,                 # Boolean determining if the target is named
    Padding(3),                            # Mystery bytes!
    'szTargetName' / Padded(255, CString('utf8')),         # Target name
                                                           # DFU File Manager does not initialise this
                                                           # memory, so our file will not exactly match
                                                           # its output.
    'dwTargetSize' / Rebuild(Int32ul, DFU_target_size),    # Total size of target images
    'dwNbElements' / Rebuild(Int32ul, len_(this.images)),  # Count the number of target images
    'images' / GreedyRange(DFU_image)
)

DFU_body = Struct(
    'szSignature' / Const(DFU_SIGNATURE),  # DFU format identifier (changes on major revisions)
    'bVersion' / Const(1, Int8ul),         # DFU format revision   (changes on minor revisions)
    'DFUImageSize' / Rebuild(Int32ul, DFU_file_length),    # Total DFU file length in bytes
    'bTargets' / Rebuild(Int8ul, len_(this.targets)),      # Number of targets in the file

    'targets' / GreedyRange(DFU_target),

    'bcdDevice' / Int16ul,                 # Firmware version, or 0xffff if ignored
    'idProduct' / Hex(Int16ul),            # USB product ID or 0xffff to ignore
    'idVendor' / Hex(Int16ul),             # USB vendor ID or 0xffff to ignore
    'bcdDFU' / Const(0x011A, Int16ul),     # DFU specification number
    'ucDfuSignature' / Const(b'UFD'),      # 0x44, 0x46 and 0x55 ie 'DFU' but reversed
    'bLength' / Const(16, Int8ul)          # Length of the DFU suffix in bytes
)

DFU = Struct(
    'fields' / RawCopy(DFU_body),
    'dwCRC' / Checksum(Int32ul,            # CRC calculated over the whole file, except for itself
                       lambda data: 0xffffffff ^ zlib.crc32(data),
                       this.fields.data)
)


def display_dfu_info(parsed):
    print(f'''
Device: {parsed.fields.value.bcdDevice}
Target: {parsed.fields.value.idProduct:04x}:{parsed.fields.value.idVendor:04x}
Size: {parsed.fields.value.DFUImageSize:,} bytes
Targets: {parsed.fields.value.bTargets}''')
    for target in parsed.fields.value.targets:
        print(f'''
    Name: {target.szTargetName}
    Alternate Setting: {target.bAlternateSetting}
    Size: {target.dwTargetSize:,} bytes
    Images: {target.dwNbElements}''')
        for image in target.images:
            print(f'''
        Offset: {image.dwElementAddress}
        Size: {len(image.data):,} bytes
''')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(help='Commands', dest='command')
    parser_build = subparsers.add_parser('build', help='Build a DFU image')
    parser_read = subparsers.add_parser('read', help='Read a DFU image')
    parser_dump = subparsers.add_parser('dump', help='Dump binary from DFU image')
    parser_dump.add_argument('--force', action='store_true')
    parser_build.add_argument('--out', type=pathlib.Path, help='Output file', default=None)
    parser_build.add_argument('--address', type=int, default=0x08000000)
    parser_build.add_argument('--force', action='store_true')
    parser.add_argument('file', type=pathlib.Path, help='Input file')
    parser.add_argument('--verbose', action='store_true')

    args = parser.parse_args()

    if not args.file.is_file():
        raise parser.error(f'Invalid input file "{args.file}"')

    if args.command == 'build':
        if not args.out.parent.is_dir():
            raise parser.error(f'Output directory "{args.out.parent}" does not exist!')
        elif args.out.is_file() and not args.force:
            raise parser.error(f'Existing output file "{args.out}", use --force to overwrite!')

        if not args.file.suffix == ".bin":
            raise parser.error(f'Input file "{args.file}", is not a .bin file?')

        output = DFU.build({'fields': {'value': {
            'targets': [{
                'bAlternateSetting': 0,
                'bTargetNamed': True,
                'szTargetName': 'ST...',
                'images': [{
                    'dwElementAddress': args.address,
                    'data': open(args.file, 'rb').read()
                }]
            }],
            'bcdDevice': 0,
            'idProduct': 0x0000,
            'idVendor': 0x0483
        }}})

        if args.verbose:
            print(f'''Packing "{args.file}" into "{args.out}"''')
            display_dfu_info(DFU.parse(output))

        open(args.out, 'wb').write(output)

    if args.command == 'read' or args.command == 'dump':
        try:
            parsed = DFU.parse(open(args.file, 'rb').read())
        except construct.core.ConstructError as error:
            parser.error(f'Invalid dfu file {args.file} ({error})')

        display_dfu_info(parsed)

        if args.command == 'dump':

            for target in parsed.fields.value.targets:
                target_id = target.bAlternateSetting
                for image in target.images:
                    address = image.dwElementAddress
                    data = image.data
                    dest = str(args.file).replace('.dfu', '')
                    filename = f"{dest}-{target_id}-{address}.bin"

                    if pathlib.Path(filename).is_file() and not args.force:
                        raise parser.error(f'Existing output file "{filename}", use --force to overwrite!')

                    print(f"Dumping image at {address} to {filename} ({len(data)} bytes)")

                    open(filename, 'wb').write(data)