#!/usr/bin/env python

# Copyright (c) 2016-2020 Satlab A/S <satlab@satlab.com>
# All rights reserved. Do not distribute without permission.

from __future__ import print_function

import argparse
import subprocess


def msg(string, **kwargs):
    print('\033[1m[+] ' + str(string) + '\033[0m', **kwargs)


def err(string, **kwargs):
    print('\033[1m[E] ' + str(string) + '\033[0m', **kwargs)


def execute(cmd):
    popen = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        universal_newlines=True)
    for stdout_line in iter(popen.stdout.readline, ""):
        yield stdout_line
    popen.stdout.close()
    return_code = popen.wait()
    if return_code:
        raise subprocess.CalledProcessError(return_code, cmd)


def main():
    VERSION = '1.2.0'

    parser = argparse.ArgumentParser('Satlab SRS-4 firmware flash tool')

    parser.add_argument('binary', help='Binary image to flash')
    parser.add_argument('--adapter-khz', default=1800, type=int,
                        help='Set programming adapter speed')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='Show verbose output')

    args = parser.parse_args()

    openocd_cmd = ['openocd',
                   '-c', 'adapter driver ftdi',
                   '-c', 'ftdi_vid_pid 0x0403 0x6010',
                   '-c', 'ftdi_layout_init 0x0c08 0x0f1b',
                   '-c', 'ftdi_layout_signal nTRST -data 0x0100 -noe 0x0400',
                   '-c', 'ftdi_layout_signal nSRST -data 0x0200 -noe 0x0800',
                   '-f', 'interface/ftdi/swd-resistor-hack.cfg',
                   '-f', 'target/stm32h7x_dual_bank.cfg',
                   '-c', 'adapter speed {}'.format(args.adapter_khz),
                   '-c', 'reset_config none',
                   '-c', 'program {} verify reset exit 0x08000000'.format(
                       args.binary)]

    msg('Satlab SRS-4 firmware flash tool v' + VERSION)
    try:
        msg('Flashing image using OpenOCD {}'.format(args.binary))
        if args.verbose:
            print(' '.join(openocd_cmd))
        openocd_out = ''
        for path in execute(openocd_cmd):
            if args.verbose:
                print(path, end='')
            openocd_out += path
    except Exception as e:
        err('Could not flash firmware, please ensure that:')
        err('  1) OpenOCD v0.10.0 or newer is installed')
        err('  2) The JTAG adapter is connected to the board')
        err('  3) You supply the correct filename for the firmware image')
        err('This error could provide a clue:')
        print(e)
        err('Output from OpenOCD:')
        print(openocd_out, end='')
    else:
        msg('Succesfully flashed the firmare to the board')


if __name__ == '__main__':
    main()
