#!/usr/bin/env python
#
# Copyright (c) 2012 Intel, Inc.
# License: GPLv2
# Author: Artem Bityutskiy <artem.bityutskiy@linux.intel.com>

# Note! We use the below docstring for the program help text as well.
"""
Write raw image file to the target block device using the block map file (AKA
bmap). The bmap contains list of blocks which have to be read from the image
file and then written to the block device. The rest of the blocks are not
required to be copied. And usually image files have a lot of useless blocks
(i.e., the blocks which are not used in the internal file-system of the image),
so flashing with bmap is usually much faster than copying entire image to the
block device.

For example, you may have a 4GiB image file, which contains only 100MiB of user
data. In this case, with the bmap file you will write only a little bit more
than 100MiB of data from the image file to the block device. This is a lot
faster than writing the entire 4GiB image. We say that it is a bit more than
100MiB because there are also file-system meta-data, partition table, etc. The
bmap fail is quite human-readable and contains a lot of commentaries. But
essentially, it is an XML document which contains list of blocks in the image
file which have to be copied to the block device.
"""

VERSION = "0.1.0"

import argparse
import os
import sys
import stat
import time
import logging
from bmaptools import BmapFlasher, BmapHelpers

def parse_arguments():
    """ A helper function which parses the input arguments. """

    parser = argparse.ArgumentParser(description = __doc__,
                                     prog = 'bmap-flasher')

    # The first positional argument - image file
    text = "the image file to flash. Supported formats: uncompressed, " + \
           ", ".join(BmapFlasher.supported_image_formats)
    parser.add_argument("image", help = text)

    # The second positional argument - block device node
    text = "the block device node to flash the image to"
    parser.add_argument("bdev", help = text)

    # The --bmap option
    text = "the block map file for the image"
    parser.add_argument("--bmap", help = text)

    # The --no-verify option
    text = "do not verify the data checksum while writing"
    parser.add_argument("--no-verify", action="store_true", help = text)

    # The --no-sync option
    text = "do not synchronize the block device after flashing (use "  \
           "carefully and make sure you synchronize the block device " \
           "manually before you unplug it)"
    parser.add_argument("--no-sync", action="store_true", help = text)

    # The --quiet option
    text = "be quiet"
    parser.add_argument("-q", "--quiet", action="store_true", help = text)

    # The --version option
    parser.add_argument("--version", action="version", \
                        version="%(prog)s " + "%s" % VERSION)

    return parser.parse_args()

def setup_logger(loglevel):
    """ A helper function which sets up and configures the logger. The log
        level is initialized to 'loglevel'. Returns the logger object. """

    # Change log level names to something less nicer than the default
    # all-capital 'INFO' etc.
    logging.addLevelName(logging.ERROR, "error!")
    logging.addLevelName(logging.WARNING, "warning!")
    logging.addLevelName(logging.DEBUG, "debug")
    logging.addLevelName(logging.INFO, "info")

    log = logging.getLogger('bmap-flasher-logger')
    log.setLevel(loglevel)
    formatter = logging.Formatter("bmap-flasher: %(levelname)s: %(message)s")
    where = logging.StreamHandler()
    where.setFormatter(formatter)
    log.addHandler(where)

    return log

def main():
    """ Script entry point. """

    args = parse_arguments()

    if args.quiet:
        log = setup_logger(logging.ERROR)
    else:
        log = setup_logger(logging.INFO)

    try:
        is_block_device = stat.S_ISBLK(os.stat(args.bdev).st_mode)
    except OSError as err:
        log.error("cannot access block device '%s': %s" \
                   % (args.bdev, err.strerror))
        raise SystemExit(1)

    if not is_block_device:
        log.warning("'%s' is not a block device!" % args.bdev)

    try:
        flasher = BmapFlasher.BmapFlasher(args.image, args.bdev, args.bmap)
    except BmapFlasher.Error as err:
        log.error(str(err))
        raise SystemExit(1)

    start_time = time.time()
    if not args.bmap:
        log.info("no block map given (see the --bmap option")
        log.info("falling-back to writing entire image to '%s'" % args.bdev)
    else:
        log.info("block map format version %s" % flasher.bmap_version)
        log.info("%d blocks of size %d (%s), mapped %d blocks (%s or %.1f%%)" \
                 % (flasher.bmap_blocks_cnt, flasher.bmap_block_size,
                    flasher.bmap_image_size_human, flasher.bmap_mapped_cnt,
                    flasher.bmap_mapped_size_human,
                    flasher.bmap_mapped_percent))
        log.info("writing the image to '%s' using bmap file '%s'" \
                 % (args.bdev, args.bmap))

    try:
        flasher.write(False, not args.no_verify)
    except BmapFlasher.Error as err:
        log.error(str(err))
        raise SystemExit(1)

    # Synchronize the block device
    if not args.no_sync:
        log.info("synchronizing block device '%s'" % args.bdev)
        try:
            flasher.sync()
        except BmapFlasher.Error as err:
            log.error(str(err))
            raise SystemExit(1)

    flashing_time = time.time() - start_time
    flashing_speed = flasher.bmap_image_size / flashing_time
    log.info("flashing time: %s, flashing speed %s/sec" \
             % (BmapHelpers.human_time(flashing_time), \
                BmapHelpers.human_size(flashing_speed)))

if __name__ == "__main__":
    sys.exit(main())
