#! /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 platform
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 BossacBinaryFlasher(ZephyrBinaryFlasher):
    '''Flasher front-end for bossac.'''

    def __init__(self, bin_name, bossac='bossac', debug=False):
        super(BossacBinaryFlasher, self).__init__(debug=debug)
        self.bin_name = bin_name
        self.bossac = bossac

    def replaces_shell_script(shell_script):
        return shell_script == 'bossa-flash.sh'

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

        Required:

        - O: build output directory
        - KERNEL_BIN_NAME: name of kernel binary

        Optional:

        - BOSSAC: path to bossac, default is bossac
        '''
        bin_name = path.join(get_env_or_bail('O'),
                             get_env_or_bail('KERNEL_BIN_NAME'))
        bossac = os.environ.get('BOSSAC', 'bossac')
        return BossacBinaryFlasher(bin_name, bossac=bossac, debug=debug)

    def flash(self, **kwargs):
        if platform.system() != 'Linux':
            msg = 'CAUTION: No flash tool for your host system found!'
            raise NotImplementedError(msg)

        cmd_stty = ['stty', '-F', '/dev/ttyACM0', 'raw', 'ispeed', '1200',
                    'ospeed', '1200', 'cs8', '-cstopb', 'ignpar', 'eol', '255',
                    'eof', '255']
        cmd_flash = [self.bossac, '-R', '-e', '-w', '-v', '-b', self.bin_name]

        check_call(cmd_stty, self.debug)
        check_call(cmd_flash, self.debug)


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.')


class Esp32BinaryFlasher(ZephyrBinaryFlasher):
    '''Flasher front-end for espidf.'''

    def __init__(self, elf, device, baud=921600, flash_size='detect',
                 flash_freq='40m', flash_mode='dio', espdif='espidf',
                 debug=False):
        super(Esp32BinaryFlasher, self).__init__(debug=debug)
        self.elf = elf
        self.device = device
        self.baud = baud
        self.flash_size = flash_size
        self.flash_freq = flash_freq
        self.flash_mode = flash_mode
        self.espdif = espdif

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

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

        Required:

        - O: build output directory
        - KERNEL_ELF_NAME: name of kernel binary in ELF format

        Optional:

        - ESP_DEVICE: serial port to flash, default /dev/ttyUSB0
        - ESP_BAUD_RATE: serial baud rate, default 921600
        - ESP_FLASH_SIZE: flash size, default 'detect'
        - ESP_FLASH_FREQ: flash frequency, default '40m'
        - ESP_FLASH_MODE: flash mode, default 'dio'
        - ESP_TOOL: complete path to espdif, or set to 'espidf' to look for it
          in $ESP_IDF_PATH/components/esptool_py/esptool/esptool.py
        '''
        elf = path.join(get_env_or_bail('O'),
                        get_env_or_bail('KERNEL_ELF_NAME'))

        # TODO add sane device defaults on other platforms than Linux.
        device = os.environ.get('ESP_DEVICE', '/dev/ttyUSB0')
        baud = os.environ.get('ESP_BAUD_RATE', '921600')
        flash_size = os.environ.get('ESP_FLASH_SIZE', 'detect')
        flash_freq = os.environ.get('ESP_FLASH_FREQ', '40m')
        flash_mode = os.environ.get('ESP_FLASH_MODE', 'dio')
        espdif = os.environ.get('ESP_TOOL', 'espidf')

        if espdif == 'espdif':
            idf_path = get_env_or_bail('ESP_IDF_PATH')
            espdif = path.join(idf_path, 'components', 'esptool_py', 'esptool',
                               'esptool.py')

        return Esp32BinaryFlasher(elf, device, baud=baud,
                                  flash_size=flash_size, flash_freq=flash_freq,
                                  flash_mode=flash_mode, espdif=espdif,
                                  debug=debug)

    def flash(self, **kwargs):
        bin_name = path.splitext(self.elf)[0] + path.extsep + 'bin'
        cmd_convert = [self.espdif, '--chip', 'esp32', 'elf2image', self.elf]
        cmd_flash = [self.espdif, '--chip', 'esp32', '--port', self.device,
                     '--baud', self.baud, '--before', 'default_reset',
                     '--after', 'hard_reset', 'write_flash', '-u',
                     '--flash_mode', self.flash_mode,
                     '--flash_freq', self.flash_freq,
                     '--flash_size', self.flash_size,
                     '0x1000', bin_name]

        print("Converting ELF to BIN")
        check_call(cmd_convert, self.debug)

        print("Flashing ESP32 on {} ({}bps)".format(self.device, self.baud))
        check_call(cmd_flash, self.debug)

class PyOcdBinaryFlasher(ZephyrBinaryFlasher):
    '''Flasher front-end for pyocd-flashtool.'''

    def __init__(self, bin_name, target, flashtool='pyocd-flashtool',
                 board_id=None, daparg=None, debug=False):
        super(PyOcdBinaryFlasher, self).__init__(debug=debug)
        self.bin_name = bin_name
        self.target = target
        self.flashtool = flashtool
        self.board_id = board_id
        self.daparg = daparg

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

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

        Required:

        - O: build output directory
        - KERNEL_BIN_NAME: name of kernel binary
        - PYOCD_TARGET: target override

        Optional:

        - PYOCD_FLASHTOOL: flash tool path, defaults to pyocd-flashtool
        - PYOCD_BOARD_ID: ID of board to flash, default is to guess
        - PYOCD_DAPARG_ARG: arguments to pass to flashtool, default is none
        '''
        bin_name = path.join(get_env_or_bail('O'),
                             get_env_or_bail('KERNEL_BIN_NAME'))
        target = get_env_or_bail('PYOCD_TARGET')

        flashtool = os.environ.get('PYOCD_FLASHTOOL', 'pyocd-flashtool')
        board_id = os.environ.get('PYOCD_BOARD_ID', None)
        daparg = os.environ.get('PYOCD_DAPARG_ARG', None)

        return PyOcdBinaryFlasher(bin_name, target,
                                  flashtool=flashtool, board_id=board_id,
                                  daparg=daparg, debug=debug)

    def flash(self, **kwargs):
        daparg_args = []
        if self.daparg is not None:
            daparg_args = ['-da', self.daparg]

        board_args = []
        if self.board_id is not None:
            board_args = ['-b', self.board_id]

        cmd = ([self.flashtool] +
               daparg_args +
               ['-t', self.target] +
               board_args +
               [self.bin_name])

        print('Flashing Target Device')
        check_call(cmd, self.debug)


# 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)