# pyOCD debugger
# Copyright (c) 2018-2020 Arm Limited
# Copyright (c) 2021 Chris Reed
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import logging
import itertools
from elftools.elf.elffile import ELFFile
from intelhex import IntelHex
import errno
from typing import (Any, Callable, Dict, IO, List, Optional, Tuple, TYPE_CHECKING, Union)

from .loader import (
    FlashLoader,
    ProgressCallback,
)
from ..core import exceptions

if TYPE_CHECKING:
    from ..core.session import Session

LOG = logging.getLogger(__name__)

def ranges(i: List[int]) -> List[Tuple[int, int]]:
    """Accepts a sorted list of byte addresses. Breaks the addresses into contiguous ranges.
    Yields 2-tuples of the start and end address for each contiguous range.

    For instance, the input [0, 1, 2, 3, 32, 33, 34, 35] will yield the following 2-tuples:
    (0, 3) and (32, 35).
    """
    for a, b in itertools.groupby(enumerate(i), lambda x: x[1] - x[0]):
        b = list(b)
        yield b[0][1], b[-1][1]

class FileProgrammer(object):
    """@brief Class to manage programming a file in any supported format with many options.

    Most specifically, this class implements the behaviour provided by the command-line flash
    programming tool. The code in this class simply extracts data from the given file, potentially
    respecting format-specific options such as the base address for binary files. Then the heavy
    lifting of flash programming is handled by FlashLoader, and beneath that, FlashBuilder.

    Support file formats are:
    - Binary (.bin)
    - Intel Hex (.hex)
    - ELF (.elf or .axf)
    """
    def __init__(self,
            session: "Session",
            progress: Optional[ProgressCallback] = None,
            chip_erase: Optional[bool] = None,
            smart_flash: Optional[bool] = None,
            trust_crc: Optional[bool] = None,
            keep_unwritten: Optional[bool] = None,
            no_reset: Optional[bool] = None
        ):
        """@brief Constructor.

        @param self
        @param session The session object.
        @param progress A progress report handler as a callable that takes a percentage completed.
            If not set or None, a default progress handler will be used unless the session option
            'hide_programming_progress' is set to True, in which case progress will be disabled.
        @param chip_erase Sets whether to use chip erase or sector erase. The value must be one of
            "auto", "sector", or "chip". "auto" means the fastest erase method should be used.
        @param smart_flash If set to True, the programmer will attempt to not program pages whose
            contents are not going to change by scanning target flash memory. A value of False will
            force all pages to be erased and programmed.
        @param trust_crc Boolean indicating whether to use only the sector CRC32 to decide whether a
            sector already contains the data to be programmed. Use with caution, as CRC32 may return
            the same value for different content.
        @param keep_unwritten Depending on the sector versus page size and the amount of data
            written, there may be ranges of flash that would be erased but not written with new
            data. This parameter sets whether the existing contents of those unwritten ranges will
            be read from memory and restored while programming.
        @param no_reset Boolean indicating whether if the device should not be reset after the
            programming process has finished.
        """
        self._session = session
        self._chip_erase = chip_erase
        self._smart_flash = smart_flash
        self._trust_crc = trust_crc
        self._keep_unwritten = keep_unwritten
        self._no_reset = no_reset
        self._progress = progress
        self._loader = None

        self._format_handlers: Dict[str, Callable[..., None]] = {
            'axf': self._program_elf,
            'bin': self._program_bin,
            'elf': self._program_elf,
            'hex': self._program_hex,
            }

    def program(self, file_or_path: Union[str, IO[bytes]], file_format: Optional[str] = None, **kwargs: Any):
        """@brief Program a file into flash.

        @param self
        @param file_or_path Either a string that is a path to a file, or a file-like object.
        @param file_format Optional file format name, one of "bin", "hex", "elf", "axf". If not provided,
            the file's extension will be used. If a file object is passed for _file_or_path_ then
            this parameter must be used to set the format.
        @param kwargs Optional keyword arguments for format-specific parameters.

        The only current format-specific keyword parameters are for the binary format:
        - `base_address`: Memory address at which to program the binary data. If not set, the base
            of the boot memory will be used.
        - `skip`: Number of bytes to skip at the start of the binary file. Does not affect the
            base address.

        @exception FileNotFoundError Provided file_or_path string does not reference a file.
        @exception ValueError Invalid argument value, for instance providing a file object but
            not setting file_format.
        """
        is_path = isinstance(file_or_path, str)

        # Check for valid path first.
        if is_path and not os.path.isfile(file_or_path): # type: ignore # (type checker doesn't use is_path)
            raise FileNotFoundError(errno.ENOENT, "No such file: '{}'".format(file_or_path))

        # If no format provided, use the file's extension.
        if not file_format:
            if is_path:
                # Extract the extension from the path.
                file_format = os.path.splitext(file_or_path)[1][1:] # type: ignore # (type checker doesn't use is_path)

                # Explicitly check for no extension.
                if file_format == '':
                    raise ValueError("file path '{}' does not have an extension and "
                                        "no format is set".format(file_or_path))
            else:
                raise ValueError("file object provided but no format is set")

        # Check the format is one we understand.
        if file_format is None or file_format not in self._format_handlers:
            raise ValueError("unknown file format '%s'" % file_format)

        self._loader = FlashLoader(self._session,
                                    progress=self._progress,
                                    chip_erase=self._chip_erase,
                                    smart_flash=self._smart_flash,
                                    trust_crc=self._trust_crc,
                                    keep_unwritten=self._keep_unwritten,
                                    no_reset=self._no_reset)

        # file_obj = None
        # Open the file if a path was provided.
        if is_path:
            mode = 'rb'
            if file_format == 'hex':
                # hex file must be read as plain text file
                mode = 'r'
            assert isinstance(file_or_path, str)
            file_obj = open(file_or_path, mode)
        else:
            assert not isinstance(file_or_path, str)
            file_obj = file_or_path
        try:

            # Pass to the format-specific programmer.
            self._format_handlers[file_format](file_obj, **kwargs)
            self._loader.commit()
        finally:
            if is_path and file_obj is not None:
                file_obj.close()

    def _program_bin(self, file_obj: IO[bytes], **kwargs: Any) -> None:
        """@brief Binary file format loader"""
        assert self._loader

        # If no base address is specified use the start of the boot memory.
        address = kwargs.get('base_address', None)
        if address is None:
            boot_memory = self._session.target.memory_map.get_boot_memory()
            if boot_memory is None:
                raise exceptions.TargetSupportError("No boot memory is defined for this device")
            address = boot_memory.start
        assert isinstance(address, int)

        skip_offset = kwargs.get('skip', 0)
        if not isinstance(skip_offset, int):
            raise TypeError("skip argument must be an integer")
        file_obj.seek(skip_offset, os.SEEK_SET)
        data = list(bytearray(file_obj.read()))

        self._loader.add_data(address, data)

    def _program_hex(self, file_obj: IO[bytes], **kwargs: Any) -> None:
        """Intel hex file format loader"""
        assert self._loader

        hexfile = IntelHex(file_obj)
        addresses = hexfile.addresses()
        addresses.sort()

        data_list = list(ranges(addresses))
        for start, end in data_list:
            size = end - start + 1
            data = list(hexfile.tobinarray(start=start, size=size))
            # Ignore invalid addresses for HEX files only
            # Binary files (obviously) don't contain addresses
            # For ELF files, any metadata that's not part of the application code
            # will be held in a section that doesn't have the SHF_WRITE flag set
            try:
                self._loader.add_data(start, data)
            except ValueError as e:
                LOG.warning("Failed to add data chunk: %s", e)

    def _program_elf(self, file_obj: IO[bytes], **kwargs: Any) -> None:
        assert self._loader

        elf = ELFFile(file_obj)
        for segment in elf.iter_segments():
            addr = segment['p_paddr']
            if segment.header.p_type == 'PT_LOAD' and segment.header.p_filesz != 0:
                data = bytearray(segment.data())
                LOG.debug("Writing segment LMA:0x%08x, VMA:0x%08x, size %d", addr,
                          segment['p_vaddr'], segment.header.p_filesz)
                try:
                    self._loader.add_data(addr, data)
                except ValueError as e:
                    LOG.warning("Failed to add data chunk: %s", e)
            else:
                LOG.debug("Skipping segment LMA:0x%08x, VMA:0x%08x, size %d", addr,
                          segment['p_vaddr'], segment.header.p_filesz)
