#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import sys
import struct
import traceback
import argparse


class MULTI_SFS_HEADER_T(object):
    MAGIC = 0x5346534D  # which means "MSFS"
    VERSION = 0
    FORMAT_STRING = "<4I16x"

    def __init__(self):
        self.magic = self.MAGIC      # 0x5346534D
        self.version = self.VERSION    # structure version
        self.fs_count = ""  # flash sequence count
        self.fs_offset = struct.calcsize(self.FORMAT_STRING)  # flash sequence offset, default is 32
        self.reserved = ""  # reserved 16 bytes


class FLASH_SEQ_INFO_T(object):
    VERSION = 0
    FORMAT_STRING = "<3IQ12x"

    def __init__(self):
        self.version = self.VERSION    # structure version
        self.offset = 0  # data offset
        self.size = 0  # data size
        self.flash_id = 0  # Memory Type (1 Byte) <<8 | Manufacturer ID (1 Byte),e.g. "gd25q32/16" -> 0x40C8
        self.reserved = ""  # reserved 16 bytes
        # image data
        self.image_data = None


class PackError(Exception):
  """Application-specific errors.

  These errors represent issues for which a stack-trace should not be
  presented.

  Attributes:
    message: Error message.
  """

  def __init__(self, message):
      self.message = message
      Exception.__init__(self, message)


class SaveImage(object):
    def __init__(self, filename):
        self.file = None
        self.filename = filename
        self._initialize()

    def _initialize(self):
        self.file = open(self.filename, 'wb+')

    def write(self, data):
        if self.file:
            self.file.write(data)
        else:
            raise Exception("Failed to create binary file.")

    def __del__(self):
        if self.file:
            self.file.flush()
            self.file.close()


class ReadImage(object):

    def __init__(self, image_filename):
        """Initializes an image handler.

        Arguments:
          image_filename: The name of the file to operate on.

        Raises:
          ValueError: If data in the file is invalid.
        """
        self.file = None
        self.image_size = 0
        self._file_pos = 0
        self.filename = image_filename
        self._initialize()

    def _initialize(self):
        try:
            self._file_pos = 0
            self.file = open(self.filename, 'rb')
            self.file.seek(0, os.SEEK_END)
            self.image_size = self.file.tell()
            self.file.seek(0, os.SEEK_SET)
        except IOError as e:
            # print("Failed to open %s" % self.filename)
            self.file = None

    def seek(self, offset):
        """Sets the cursor position for reading from unsparsified file.

        Arguments:
          offset: Offset to seek to from the beginning of the file.
        """
        self._file_pos = offset

    def read(self, size):
        """Reads data from the unsparsified file.

        This method may return fewer than |size| bytes of data if the end
        of the file was encountered.

        The file cursor for reading is advanced by the number of bytes
        read.

        Arguments:
          size: Number of bytes to read.

        Returns:
          The data.

        """
        self.file.seek(self._file_pos)
        data = self.file.read(size)
        self._file_pos += len(data)
        return data

    def tell(self):
        """Returns the file cursor position for reading from unsparsified file.

        Returns:
          The file cursor position for reading.
        """
        return self._file_pos


class Packet(object):
    """Business logic for msfstool command-line tool."""

    def make_msfs_image(self, output, images):
        """Implementation of the 'make_msfs_image' command.

        Arguments:
          output: Output file where disk image is to be written to.
          images: List of partition image paths to be combined (as specified by
                  bpt).  Each element is of the form.
                  'FLASH_ID@/PATH/TO/SFS_IMAGE'

        Raises:
          PackError: If application-specific error occurs.
        """
        # Create mapping of partition name to partition image file.
        image_file_names = {}
        try:
            if images:
                for item in images:
                    flash_id, path = item.split("@")
                    flash_id = flash_id.lower()
                    if 0 == flash_id.find("0x"):
                        id = int(flash_id, 16)
                    else:
                        id = int(flash_id, 10)
                    image_file_names[id] = path
                    if not os.path.exists(path):
                        raise PackError('image file do not exist: {}.'.format(item))
        except ValueError as e:
            raise PackError('Bad image argument {}.'.format(images))

        _ids = []
        for k, v in image_file_names.items():
            print("Flash sequence info :0x%X : %s" % (k, v))
            _ids.append(k)
        _ids.sort(reverse=False)
        _fs_infos = []
        for k in _ids:
            _image = ReadImage(image_file_names[k])
            if _image.file is None:
                raise PackError('Failed to open sfs fie: {}.'.format(image_file_names[k]))
            _info = FLASH_SEQ_INFO_T()
            _info.version = FLASH_SEQ_INFO_T.VERSION
            _info.size = _image.image_size
            _info.flash_id = k
            _info.image_data = _image.read(_image.image_size)
            _fs_infos.append(_info)

        # make image
        _header = MULTI_SFS_HEADER_T()
        _header.fs_count = len(_ids)
        _header.fs_offset = struct.calcsize(MULTI_SFS_HEADER_T.FORMAT_STRING)
        _header_data = struct.pack(MULTI_SFS_HEADER_T.FORMAT_STRING,
                                   _header.magic,
                                   _header.version,
                                   _header.fs_count,
                                   _header.fs_offset)

        # create output image file
        _s_file = SaveImage(output)
        _s_file.write(_header_data)

        # info
        offset = struct.calcsize(MULTI_SFS_HEADER_T.FORMAT_STRING) + \
                 _header.fs_count*struct.calcsize(FLASH_SEQ_INFO_T.FORMAT_STRING)
        for item in _fs_infos:
            item.offset = offset
            _info_data = struct.pack(FLASH_SEQ_INFO_T.FORMAT_STRING,
                                     item.version,
                                     item.offset,
                                     item.size,
                                     item.flash_id)
            _s_file.write(_info_data)
            offset += item.size

        for item in _fs_infos:
            _s_file.write(item.image_data)
        return True

    def extract_msfs(self, output_path, image_file):
        output_path = output_path.replace('\\', '/')
        os.makedirs(output_path, exist_ok=True)
        _image = ReadImage(image_file)
        if _image.file is None:
            raise PackError('Failed to open msfs fie: {}.'.format(image_file))
        if _image.image_size < struct.calcsize(MULTI_SFS_HEADER_T.FORMAT_STRING):
            raise PackError('Invalid size of the msfs fie: {}.'.format(image_file))
        _header_bin = _image.read(struct.calcsize(MULTI_SFS_HEADER_T.FORMAT_STRING))
        _header = MULTI_SFS_HEADER_T()
        (_header.magic,
         _header.version,
         _header.fs_count,
         _header.fs_offset) = struct.unpack(MULTI_SFS_HEADER_T.FORMAT_STRING, _header_bin)
        if not _header.magic == MULTI_SFS_HEADER_T.MAGIC:
            raise PackError('Invalid msfs fie: {}.'.format(image_file))
        print("msfs summary information :")
        print("     magic: 0x%X" % _header.magic)
        print("   version: 0x%X" % _header.version)
        print("  fs count: 0x%X" % _header.fs_count)
        print(" fs offset: 0x%X" % _header.fs_offset)

        _fs_infos = []
        for i in range(_header.fs_count):
            _info_bin = _image.read(struct.calcsize(FLASH_SEQ_INFO_T.FORMAT_STRING))
            _info = FLASH_SEQ_INFO_T()
            (_info.version,
             _info.offset,
             _info.size,
             _info.flash_id) = struct.unpack(FLASH_SEQ_INFO_T.FORMAT_STRING, _info_bin)
            _fs_infos.append(_info)

        for item in _fs_infos:
            _sfs = "%s/FlashID_0x%X.img" % (output_path, item.flash_id)
            _file = SaveImage(_sfs)
            _image.seek(item.offset)
            _file.write(_image.read(item.size))

            print("\nextracted sfs information : %s" % _sfs)
            print("  version: 0x%X" % item.version)
            print("   offset: 0x%X" % item.offset)
            print("     size: 0x%X" % item.size)
            print(" flash id: 0x%X" % item.flash_id)

        return True


class MsfsTool(object):
    """Object for command-line tool."""

    def __init__(self):
        """Initializer method."""
        self.packet = Packet()
        pass

    def version(self, _):
        """Implements the 'version' sub-command."""
        print('1.21.32.1')

    def make_msfs(self, args):
        """Implements the 'make_msfs' sub-command."""
        if 0 == len(args.output):
            print('Option --ouptut parameter is required.')
            return False
        result = False
        try:
            result = self.packet.make_msfs_image(args.output, args.image)
        except PackError as e:
            print('Error : {}'.format(e.message))
        finally:
            if result:
                print('Successfully generated msfs file.')
            else:
                os.remove(args.output)
                print('Failed to generate msfs file.')
            return result

    def extract_msfs(self, args):
        """Implements the 'extract_msfs' sub-command."""
        result = False
        try:
            result = self.packet.extract_msfs(args.output, args.image)
        except PackError as e:
            print('Error : {}'.format(e.message))
        finally:
            if result:
                print('Successfully extracted msfs file.')
            else:
                print('Failed to extract msfs file.')
            return result

    def run(self, argv):
        parser = argparse.ArgumentParser()
        subparsers = parser.add_subparsers(title='subcommands')
        sub_parser = subparsers.add_parser('version', help='Prints version of MsfsTool.')
        sub_parser.set_defaults(func=self.version)

        sub_parser = subparsers.add_parser('make_msfs', help='Create multi sfs image file.')
        sub_parser.add_argument('--output',
                                help='Path to image output.',
                                type=str,
                                required=True)
        sub_parser.add_argument('--image',
                                help='Flash ID and path to sfs/nfs image file.',
                                metavar='FLASH_ID@SFS_FILE',
                                action='append',
                                required=True)
        sub_parser.set_defaults(func=self.make_msfs)

        sub_parser = subparsers.add_parser('extract_msfs', help='Extract multi sfs image file.')
        sub_parser.add_argument('--output',
                                help='Path to extract sfs.',
                                type=str,
                                required=True)
        sub_parser.add_argument('--image',
                                help='Path to msfs file.',
                                type=str,
                                required=True)
        sub_parser.set_defaults(func=self.extract_msfs)

        args = parser.parse_args(argv[1:])
        try:
            func = args.func
        except AttributeError:
            parser.error("invalid arguments")
        return func(args)


if __name__ == '__main__':
    obj = MsfsTool()
    obj.run(sys.argv)