# -*- coding: utf-8 -*-
"""
This module contains functions for turning a Python script into a .hex file
and flashing it onto a BBC micro:bit.

Copyright (c) 2015-2018 Nicholas H.Tollervey and others.

See the LICENSE file for more information, or visit:

https://opensource.org/licenses/MIT
"""

from __future__ import print_function

import argparse
import binascii
import ctypes
import os
import struct
import sys
from subprocess import check_output
import time

# nudatus is an optional dependancy
can_minify = True
try:
    import nudatus
except ImportError:  # pragma: no cover
    can_minify = False

# The default Debian MicroPython runtime is provided in the
# firmware-microbit-micropython package
DEFAULT_RUNTIME_PATH = "/usr/share/firmware-microbit-micropython/firmware.hex"

#: The magic start address in flash memory for a Python script.
_SCRIPT_ADDR = 0x3e000


#: The help text to be shown when requested.
_HELP_TEXT = """
Flash Python onto the BBC micro:bit or extract Python from a .hex file.

If no path to the micro:bit is provided uflash will attempt to autodetect the
correct path to the device. If no path to the Python script is provided uflash
will flash the unmodified MicroPython firmware onto the device. Use the -e flag
to recover a Python script from a hex file. Use the -r flag to specify a custom
version of the MicroPython runtime.

Documentation is here: https://uflash.readthedocs.io/en/latest/
"""


#: MAJOR, MINOR, RELEASE, STATUS [alpha, beta, final], VERSION of uflash
_VERSION = (1, 2, 4, )
_MAX_SIZE = 8188


#: The version number reported by the bundled MicroPython in os.uname().
MICROPYTHON_VERSION = '1.0.1'


def get_version():
    """
    Returns a string representation of the version information of this project.
    """
    return '.'.join([str(i) for i in _VERSION])


def get_minifier():
    """
    Report the minifier will be used when minify=True
    """
    if can_minify:
        return 'nudatus'
    return None


def strfunc(raw):
    """
    Compatibility for 2 & 3 str()
    """
    return str(raw) if sys.version_info[0] == 2 else str(raw, 'utf-8')


def hexlify(script, minify=False):
    """
    Takes the byte content of a Python script and returns a hex encoded
    version of it.

    Based on the hexlify script in the microbit-micropython repository.
    """
    if not script:
        return ''
    # Convert line endings in case the file was created on Windows.
    script = script.replace(b'\r\n', b'\n')
    script = script.replace(b'\r', b'\n')
    if minify:
        if not can_minify:
            raise ValueError("No minifier is available")
        script = nudatus.mangle(script.decode('utf-8')).encode('utf-8')
    # Add header, pad to multiple of 16 bytes.
    data = b'MP' + struct.pack('<H', len(script)) + script
    # Padding with null bytes in a 2/3 compatible way
    data = data + (b'\x00' * (16 - len(data) % 16))
    if len(data) > _MAX_SIZE:
        # 'MP' = 2 bytes, script length is another 2 bytes.
        raise ValueError("Python script must be less than 8188 bytes.")
    # Convert to .hex format.
    output = [':020000040003F7']  # extended linear address, 0x0003.
    addr = _SCRIPT_ADDR
    for i in range(0, len(data), 16):
        chunk = data[i:min(i + 16, len(data))]
        chunk = struct.pack('>BHB', len(chunk), addr & 0xffff, 0) + chunk
        checksum = (-(sum(bytearray(chunk)))) & 0xff
        hexline = ':%s%02X' % (strfunc(binascii.hexlify(chunk)).upper(),
                               checksum)
        output.append(hexline)
        addr += 16
    return '\n'.join(output)


def unhexlify(blob):
    """
    Takes a hexlified script and turns it back into a string of Python code.
    """
    lines = blob.split('\n')[1:]
    output = []
    for line in lines:
        # Discard the address, length etc. and reverse the hexlification
        output.append(binascii.unhexlify(line[9:-2]))
    # Check the header is correct ("MP<size>")
    if (output[0][0:2].decode('utf-8') != u'MP'):
        return ''
    # Strip off header
    output[0] = output[0][4:]
    # and strip any null bytes from the end
    output[-1] = output[-1].strip(b'\x00')
    script = b''.join(output)
    try:
        result = script.decode('utf-8')
        return result
    except UnicodeDecodeError:
        # Return an empty string because in certain rare circumstances (where
        # the source hex doesn't include any embedded Python code) this
        # function may be passed in "raw" bytes from MicroPython.
        return ''


def embed_hex(runtime_hex, python_hex=None):
    """
    Given a string representing the MicroPython runtime hex, will embed a
    string representing a hex encoded Python script into it.

    Returns a string representation of the resulting combination.

    Will raise a ValueError if the runtime_hex is missing.

    If the python_hex is missing, it will return the unmodified runtime_hex.
    """
    if not runtime_hex:
        raise ValueError('MicroPython runtime hex required.')
    if not python_hex:
        return runtime_hex
    py_list = python_hex.split()
    runtime_list = runtime_hex.split()
    embedded_list = []
    # The embedded list should be the original runtime with the Python based
    # hex embedded two lines from the end.
    embedded_list.extend(runtime_list[:-5])
    embedded_list.extend(py_list)
    embedded_list.extend(runtime_list[-5:])
    return '\n'.join(embedded_list) + '\n'


def extract_script(embedded_hex):
    """
    Given a hex file containing the MicroPython runtime and an embedded Python
    script, will extract the original Python script.

    Returns a string containing the original embedded script.
    """
    hex_lines = embedded_hex.split('\n')
    script_addr_high = hex((_SCRIPT_ADDR >> 16) & 0xffff)[2:].upper().zfill(4)
    script_addr_low = hex(_SCRIPT_ADDR & 0xffff)[2:].upper().zfill(4)
    start_script = None
    within_range = False
    # Look for the script start address
    for loc, val in enumerate(hex_lines):
        if val[0:9] == ':02000004':
            # Reached an extended address record, check if within script range
            within_range = val[9:13].upper() == script_addr_high
        elif within_range and val[0:3] == ':10' and \
                val[3:7].upper() == script_addr_low:
            start_script = loc
            break
    if start_script:
        # Find the end of the script
        end_script = None
        for loc, val in enumerate(hex_lines[start_script:]):
            if val[9:41] == 'FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF':
                end_script = loc + start_script
                break
        # Pass the extracted hex through unhexlify
        return unhexlify('\n'.join(
            hex_lines[start_script - 1:end_script if end_script else -6]))
    return ''


def find_microbit():
    """
    Returns a path on the filesystem that represents the plugged in BBC
    micro:bit that is to be flashed. If no micro:bit is found, it returns
    None.

    Works on Linux, OSX and Windows. Will raise a NotImplementedError
    exception if run on any other operating system.
    """
    # Check what sort of operating system we're on.
    if os.name == 'posix':
        # 'posix' means we're on Linux or OSX (Mac).
        # Call the unix "mount" command to list the mounted volumes.
        mount_output = check_output('mount').splitlines()
        mounted_volumes = [x.split()[2] for x in mount_output]
        for volume in mounted_volumes:
            if volume.endswith(b'MICROBIT'):
                return volume.decode('utf-8')  # Return a string not bytes.
    elif os.name == 'nt':
        # 'nt' means we're on Windows.

        def get_volume_name(disk_name):
            """
            Each disk or external device connected to windows has an attribute
            called "volume name". This function returns the volume name for
            the given disk/device.

            Code from http://stackoverflow.com/a/12056414
            """
            vol_name_buf = ctypes.create_unicode_buffer(1024)
            ctypes.windll.kernel32.GetVolumeInformationW(
                ctypes.c_wchar_p(disk_name), vol_name_buf,
                ctypes.sizeof(vol_name_buf), None, None, None, None, 0)
            return vol_name_buf.value

        #
        # In certain circumstances, volumes are allocated to USB
        # storage devices which cause a Windows popup to raise if their
        # volume contains no media. Wrapping the check in SetErrorMode
        # with SEM_FAILCRITICALERRORS (1) prevents this popup.
        #
        old_mode = ctypes.windll.kernel32.SetErrorMode(1)
        try:
            for disk in 'ABCDEFGHIJKLMNOPQRSTUVWXYZ':
                path = '{}:\\'.format(disk)
                #
                # Don't bother looking if the drive isn't removable
                #
                if ctypes.windll.kernel32.GetDriveTypeW(path) != 2:
                    continue
                if os.path.exists(path) and \
                        get_volume_name(path) == 'MICROBIT':
                    return path
        finally:
            ctypes.windll.kernel32.SetErrorMode(old_mode)
    else:
        # No support for unknown operating systems.
        raise NotImplementedError('OS "{}" not supported.'.format(os.name))


def save_hex(hex_file, path):
    """
    Given a string representation of a hex file, this function copies it to
    the specified path thus causing the device mounted at that point to be
    flashed.

    If the hex_file is empty it will raise a ValueError.

    If the filename at the end of the path does not end in '.hex' it will raise
    a ValueError.
    """
    if not hex_file:
        raise ValueError('Cannot flash an empty .hex file.')
    if not path.endswith('.hex'):
        raise ValueError('The path to flash must be for a .hex file.')
    with open(path, 'wb') as output:
        output.write(hex_file.encode('ascii'))


def flash(path_to_python=None, paths_to_microbits=None,
          path_to_runtime=None, python_script=None, minify=False):
    """
    Given a path to or source of a Python file will attempt to create a hex
    file and then flash it onto the referenced BBC micro:bit.

    If the path_to_python & python_script are unspecified it will simply flash
    the unmodified MicroPython runtime onto the device.

    If used, the python_script argument should be a bytes object representing
    a UTF-8 encoded string. For example::

        script = "from microbit import *\\ndisplay.scroll('Hello, World!')"
        uflash.flash(python_script=script.encode('utf-8'))

    If paths_to_microbits is unspecified it will attempt to find the device's
    path on the filesystem automatically.

    If the path_to_runtime is unspecified it will use the built in version of
    the MicroPython runtime. This feature is useful if a custom build of
    MicroPython is available.

    If the automatic discovery fails, then it will raise an IOError.
    """
    # Check for the correct version of Python.
    if not ((sys.version_info[0] == 3 and sys.version_info[1] >= 3) or
            (sys.version_info[0] == 2 and sys.version_info[1] >= 7)):
        raise RuntimeError('Will only run on Python 2.7, or 3.3 and later.')
    # Grab the Python script (if needed).
    python_hex = ''
    if path_to_python:
        if not path_to_python.endswith('.py'):
            raise ValueError('Python files must end in ".py".')
        with open(path_to_python, 'rb') as python_script:
            python_hex = hexlify(python_script.read(), minify)
    elif python_script:
        python_hex = hexlify(python_script, minify)

    runtime = _RUNTIME
    # Load the hex for the runtime.
    if path_to_runtime:
        with open(path_to_runtime) as runtime_file:
            runtime = runtime_file.read()
    # Generate the resulting hex file.
    micropython_hex = embed_hex(runtime, python_hex)
    # Find the micro:bit.
    if not paths_to_microbits:
        found_microbit = find_microbit()
        if found_microbit:
            paths_to_microbits = [found_microbit]
    # Attempt to write the hex file to the micro:bit.
    if paths_to_microbits:
        for path in paths_to_microbits:
            hex_path = os.path.join(path, 'micropython.hex')
            print('Flashing Python to: {}'.format(hex_path))
            save_hex(micropython_hex, hex_path)
    else:
        raise IOError('Unable to find micro:bit. Is it plugged in?')


def extract(path_to_hex, output_path=None):
    """
    Given a path_to_hex file this function will attempt to extract the
    embedded script from it and save it either to output_path or stdout
    """
    with open(path_to_hex, 'r') as hex_file:
        python_script = extract_script(hex_file.read())
        if output_path:
            with open(output_path, 'w') as output_file:
                output_file.write(python_script)
        else:
            print(python_script)


def watch_file(path, func, *args, **kwargs):
    """
    Watch a file for changes by polling its last modification time. Call the
    provided function with *args and **kwargs upon modification.
    """
    if not path:
        raise ValueError('Please specify a file to watch')
    print('Watching "{}" for changes'.format(path))
    last_modification_time = os.path.getmtime(path)
    try:
        while True:
            time.sleep(1)
            new_modification_time = os.path.getmtime(path)
            if new_modification_time == last_modification_time:
                continue
            func(*args, **kwargs)
            last_modification_time = new_modification_time
    except KeyboardInterrupt:
        pass


def main(argv=None):
    """
    Entry point for the command line tool 'uflash'.

    Will print help text if the optional first argument is "help". Otherwise
    it will ensure the optional first argument ends in ".py" (the source
    Python script).

    An optional second argument is used to reference the path to the micro:bit
    device. Any more arguments are ignored.

    Exceptions are caught and printed for the user.
    """
    if not argv:
        argv = sys.argv[1:]

    parser = argparse.ArgumentParser(description=_HELP_TEXT)
    parser.add_argument('source', nargs='?', default=None)
    parser.add_argument('target', nargs='*', default=None)
    parser.add_argument('-r', '--runtime', default=None,
                        help="Use the referenced MicroPython runtime.")
    parser.add_argument('-e', '--extract',
                        action='store_true',
                        help=("Extract python source from a hex file"
                              " instead of creating the hex file."), )
    parser.add_argument('-w', '--watch',
                        action='store_true',
                        help='Watch the source file for changes.')
    parser.add_argument('-m', '--minify',
                        action='store_true',
                        help='Minify the source')
    parser.add_argument('--version', action='version',
                        version='%(prog)s ' + get_version())
    args = parser.parse_args(argv)

    if args.extract:
        try:
            extract(args.source, args.target)
        except Exception as ex:
            error_message = "Error extracting {source}: {error!s}"
            print(error_message.format(source=args.source, error=ex),
                  file=sys.stderr)
            sys.exit(1)

    elif args.watch:
        try:
            watch_file(args.source, flash,
                       path_to_python=args.source,
                       paths_to_microbits=args.target,
                       path_to_runtime=args.runtime)
        except Exception as ex:
            error_message = "Error watching {source}: {error!s}"
            print(error_message.format(source=args.source, error=ex),
                  file=sys.stderr)
            sys.exit(1)

    else:
        try:
            flash(path_to_python=args.source, paths_to_microbits=args.target,
                  path_to_runtime=args.runtime, minify=args.minify)
        except Exception as ex:
            error_message = (
                "Error flashing {source} to {target}{runtime}: {error!s}"
            )
            source = args.source
            target = args.target if args.target else "microbit"
            if args.runtime:
                runtime = "with runtime {runtime}".format(runtime=args.runtime)
            else:
                runtime = ""
            print(error_message.format(source=source, target=target,
                                       runtime=runtime, error=ex),
                  file=sys.stderr)
            sys.exit(1)


#: A string representation of the MicroPython runtime hex.
_RUNTIME = open(DEFAULT_RUNTIME_PATH).read()


if __name__ == '__main__':  # pragma: no cover
    main(sys.argv[1:])