#! /usr/bin/env python3

# Copyright (c) 2017 Linaro Limited.
#
# SPDX-License-Identifier: Apache-2.0

"""Zephyr board flashing script

This script is a transparent replacement for an existing Zephyr flash
script. If it can invoke the flashing tools natively, it will do so; otherwise,
it delegates to the shell script passed as second argument."""

import abc
from os import path
import os
import pprint
import sys
import subprocess
import time


def get_env_or_bail(env_var):
    try:
        return os.environ[env_var]
    except KeyError:
        print('Variable {} not in environment:'.format(
                  env_var), file=sys.stderr)
        pprint.pprint(dict(os.environ), stream=sys.stderr)
        raise


def get_env_bool_or(env_var, default_value):
    try:
        return bool(int(os.environ[env_var]))
    except KeyError:
        return default_value


def check_call(cmd, debug):
    if debug:
        print(' '.join(cmd))
    subprocess.check_call(cmd)


def check_output(cmd, debug):
    if debug:
        print(' '.join(cmd))
    return subprocess.check_output(cmd)


class ZephyrBinaryFlasher(abc.ABC):
    '''Abstract superclass for flasher objects.'''

    def __init__(self, debug=False):
        self.debug = debug

    @staticmethod
    def create_for_shell_script(shell_script, debug):
        '''Factory for using as a drop-in replacement to a shell script.

        Get flasher instance to use in place of shell_script, deriving
        flasher configuration from the environment.'''
        for sub_cls in ZephyrBinaryFlasher.__subclasses__():
            if sub_cls.replaces_shell_script(shell_script):
                return sub_cls.create_from_env(debug)
        raise ValueError('no flasher replaces script {}'.format(shell_script))

    @staticmethod
    @abc.abstractmethod
    def replaces_shell_script(shell_script):
        '''Check if this flasher class replaces FLASH_SCRIPT=shell_script.'''

    @staticmethod
    @abc.abstractmethod
    def create_from_env(debug):
        '''Create new flasher instance from environment variables.

        This class must be able to replace the current FLASH_SCRIPT. The
        environment variables expected by that script are used to build
        the flasher in a backwards-compatible manner.'''

    @abc.abstractmethod
    def flash(self, **kwargs):
        '''Flash the board.'''


class DfuUtilBinaryFlasher(ZephyrBinaryFlasher):
    '''Flasher front-end for dfu-util.'''

    def __init__(self, pid, alt, img, dfuse=None, exe='dfu-util', debug=False):
        super(DfuUtilBinaryFlasher, self).__init__(debug=debug)
        self.pid = pid
        self.alt = alt
        self.img = img
        self.dfuse = dfuse
        self.exe = exe
        try:
            self.list_pattern = ', alt={}'.format(int(self.alt))
        except ValueError:
            self.list_pattern = ', name={}'.format(self.alt)

    def replaces_shell_script(shell_script):
        return shell_script == 'dfuutil.sh'

    def create_from_env(debug):
        '''Create flasher from environment.

        Required:

        - DFUUTIL_PID: USB VID:PID of the board
        - DFUUTIL_ALT: interface alternate setting number or name
        - DFUUTIL_IMG: binary to flash

        Optional:

        - DFUUTIL_DFUSE_ADDR: target address if the board is a
          DfuSe device. Ignored if not present.
        - DFUUTIL: dfu-util executable, defaults to dfu-util.
        '''
        pid = get_env_or_bail('DFUUTIL_PID')
        alt = get_env_or_bail('DFUUTIL_ALT')
        img = get_env_or_bail('DFUUTIL_IMG')
        dfuse = os.environ.get('DFUUTIL_DFUSE_ADDR', None)
        exe = os.environ.get('DFUUTIL', 'dfu-util')

        return DfuUtilBinaryFlasher(pid, alt, img, dfuse=dfuse, exe=exe,
                                    debug=debug)

    def find_device(self):
        output = check_output([self.exe, '-l'], self.debug)
        output = output.decode(sys.getdefaultencoding())
        return self.list_pattern in output

    def flash(self, **kwargs):
        reset = 0
        if not self.find_device():
            reset = 1
            print('Please reset your board to switch to DFU mode...')
            while not self.find_device():
                time.sleep(0.1)

        cmd = [self.exe, '-d,{}'.format(self.pid)]
        if self.dfuse is not None:
            cmd.extend(['-s', '{}:leave'.format(self.dfuse)])
        cmd.extend(['-a', self.alt, '-D', self.img])
        check_call(cmd, self.debug)
        if reset:
            print('Now reset your board again to switch back to runtime mode.')


# TODO: Stop using environment variables.
#
# Migrate the build system so we can use an argparse.ArgumentParser and
# per-flasher subparsers, so invoking the script becomes something like:
#
#   python zephyr_flash_debug.py openocd --openocd-bin=/openocd/path ...
#
# For now, maintain compatibility.
def flash(shell_script_full, debug):
    shell_script = path.basename(shell_script_full)
    try:
        flasher = ZephyrBinaryFlasher.create_for_shell_script(shell_script,
                                                              debug)
    except ValueError:
        # Can't create a flasher; fall back on shell script.
        check_call([shell_script_full, 'flash'], debug)
        return

    flasher.flash()


if __name__ == '__main__':
    debug = True
    try:
        debug = get_env_bool_or('KBUILD_VERBOSE', False)
        if len(sys.argv) != 3 or sys.argv[1] != 'flash':
            raise ValueError('usage: {} flash path-to-script'.format(
                sys.argv[0]))
        flash(sys.argv[2], debug)
    except Exception as e:
        if debug:
            raise
        else:
            print('Error: {}'.format(e), file=sys.stderr)
            print('Re-run with KBUILD_VERBOSE=1 for a stack trace.',
                  file=sys.stderr)
            sys.exit(1)