#!/usr/bin/env python3

# Copyright 2020 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Flashes firmware using Segger J-Link.

This script requires Segger hardware attached via JTAG/SWD.

See
https://chromium.googlesource.com/chromiumos/platform/ec/+/HEAD/docs/fingerprint/fingerprint-debugging.md#flash
for instructions.
"""

import argparse
import logging
import os
import shutil
import socket
import subprocess
import sys
import tempfile
import time


DEFAULT_SEGGER_REMOTE_PORT = 19020

# Commands are documented here: https://wiki.segger.com/J-Link_Commander
JLINK_COMMANDS = '''
exitonerror 1
r
loadfile {FIRMWARE} {FLASH_ADDRESS}
r
go
exit
'''


class BoardConfig:
    def __init__(self, interface, device, flash_address):
        self.interface = interface
        self.device = device
        self.flash_address = flash_address


SWD_INTERFACE = 'SWD'
STM32_DEFAULT_FLASH_ADDRESS = '0x8000000'
DRAGONCLAW_CONFIG = BoardConfig(interface=SWD_INTERFACE, device='STM32F412CG',
                                flash_address=STM32_DEFAULT_FLASH_ADDRESS)
ICETOWER_CONFIG = BoardConfig(interface=SWD_INTERFACE, device='STM32H743ZI',
                              flash_address=STM32_DEFAULT_FLASH_ADDRESS)

BOARD_CONFIGS = {
    'dragonclaw': DRAGONCLAW_CONFIG,
    'bloonchipper': DRAGONCLAW_CONFIG,
    'nucleo-f412zg': DRAGONCLAW_CONFIG,
    'dartmonkey': ICETOWER_CONFIG,
    'icetower': ICETOWER_CONFIG,
    'nucleo-dartmonkey': ICETOWER_CONFIG,
    'nucleo-h743zi': ICETOWER_CONFIG,
}


def is_tcp_port_open(host: str, tcp_port: int) -> bool:
    """Checks if the TCP host port is open."""

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(2)  # 2 Second Timeout
    try:
        sock.connect((host, tcp_port))
        sock.shutdown(socket.SHUT_RDWR)
    except ConnectionRefusedError:
        return False
    except socket.timeout:
        return False
    finally:
        sock.close()
    # Other errors are propagated as odd exceptions.

    # We shutdown and closed the connection, but the server may need a second
    # to start listening again. If the following error is seen, this timeout
    # should be increased. 300ms seems to be the minimum.
    #
    # Connecting to J-Link via IP...FAILED: Can not connect to J-Link via \
    #   TCP/IP (127.0.0.1, port 19020)
    time.sleep(0.5)
    return True


def create_jlink_command_file(firmware_file, config):
    tmp = tempfile.NamedTemporaryFile()
    tmp.write(JLINK_COMMANDS.format(FIRMWARE=firmware_file,
                                    FLASH_ADDRESS=config.flash_address).encode(
        'utf-8'))
    tmp.flush()
    return tmp


def flash(jlink_exe, remote, device, interface, cmd_file):
    cmd = [
        jlink_exe,
    ]

    if remote:
        logging.debug(f'Connecting to J-Link over TCP/IP {remote}.')
        remote_components = remote.split(':')
        if len(remote_components) not in [1, 2]:
            logging.debug(f'Given remote "{remote}" is malformed.')
            return 1

        host = remote_components[0]
        try:
            ip = socket.gethostbyname(host)
        except socket.gaierror as e:
            logging.error(f'Failed to resolve host "{host}": {e}.')
            return 1
        logging.debug(f'Resolved {host} as {ip}.')
        port = DEFAULT_SEGGER_REMOTE_PORT

        if len(remote_components) == 2:
            try:
                port = int(remote_components[1])
            except ValueError:
                logging.error(
                    f'Given remote port "{remote_components[1]}" is malformed.')
                return 1

        remote = f'{ip}:{port}'

        logging.debug(f'Checking connection to {remote}.')
        if not is_tcp_port_open(ip, port):
            logging.error(
                f'JLink server doesn\'t seem to be listening on {remote}.')
            logging.error('Ensure that JLinkRemoteServerCLExe is running.')
            return 1
        cmd.extend(['-ip', remote])

    cmd.extend([
        '-device', device,
        '-if', interface,
        '-speed', 'auto',
        '-autoconnect', '1',
        '-CommandFile', cmd_file,
        ])
    logging.debug('Running command: "%s"', ' '.join(cmd))
    completed_process = subprocess.run(cmd)
    logging.debug('JLink return code: %d', completed_process.returncode)
    return completed_process.returncode


def main(argv: list):

    parser = argparse.ArgumentParser()

    default_jlink = './JLink_Linux_V684a_x86_64/JLinkExe'
    if shutil.which(default_jlink) is None:
        default_jlink = 'JLinkExe'
    parser.add_argument(
        '--jlink', '-j',
        help='JLinkExe path (default: ' + default_jlink + ')',
        default=default_jlink)

    parser.add_argument(
        '--remote', '-n',
        help='Use TCP/IP host[:port] to connect to a J-Link or '
        'JLinkRemoteServerCLExe. If unspecified, connect over USB.')

    default_board = 'bloonchipper'
    parser.add_argument(
        '--board', '-b',
        help='Board (default: ' + default_board + ')',
        default=default_board)

    default_firmware = os.path.join('./build', default_board, 'ec.bin')
    parser.add_argument(
        '--image', '-i',
        help='Firmware binary (default: ' + default_firmware + ')',
        default=default_firmware)

    log_level_choices = ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL']
    parser.add_argument(
        '--log_level', '-l',
        choices=log_level_choices,
        default='DEBUG'
    )

    args = parser.parse_args(argv)
    logging.basicConfig(level=args.log_level)

    if args.board not in BOARD_CONFIGS:
        logging.error('Unable to find a config for board: "%s"', args.board)
        sys.exit(1)

    config = BOARD_CONFIGS[args.board]

    args.image = os.path.realpath(args.image)
    args.jlink = args.jlink

    cmd_file = create_jlink_command_file(args.image, config)
    ret_code = flash(args.jlink, args.remote, config.device, config.interface,
                     cmd_file.name)
    cmd_file.close()
    return ret_code


if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))