#!/usr/bin/env python
# Copyright (C) 2011 Linaro
#
# Author: Jeremy Chang <jeremy.chang@linaro.org>
#
# This file is part of Linaro Image Tools.
#
# Linaro Image Tools is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Linaro Image Tools is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Linaro Image Tools.  If not, see <http://www.gnu.org/licenses/>.

import atexit
import os
import sys
import tempfile

from linaro_image_tools import cmd_runner

from linaro_image_tools.media_create.android_boards import (
    get_board_config,
    )
from linaro_image_tools.media_create.check_device import (
    confirm_device_selection_and_ensure_it_is_ready)
from linaro_image_tools.media_create.partitions import (
    Media,
    setup_android_partitions,
    partition_mounted,
    )
from linaro_image_tools.media_create.rootfs import populate_partition
from linaro_image_tools.media_create.unpack_binary_tarball import (
    unpack_android_binary_tarball
    )
from linaro_image_tools.media_create import get_android_args_parser
from linaro_image_tools.utils import (
    additional_android_option_checks,
    ensure_command,
    get_logger
    )


# Just define the global variables
TMP_DIR = None
BOOT_DISK = None
SYSTEM_DISK = None
CACHE_DISK = None
DATA_DISK = None
SDCARD_DISK = None


# Registered as the first atexit handler as we want this to be the last
# handler to execute.
@atexit.register
def cleanup_tempdir():
    """Remove TEMP_DIR with all its contents.

    Before doing so, make sure DISKs are not mounted.
    """
    devnull = open('/dev/null', 'w')
    # ignore non-zero return codes
    for disk in BOOT_DISK, SYSTEM_DISK, CACHE_DISK, DATA_DISK, \
                SDCARD_DISK:
        if disk is not None:
            try:
                cmd_runner.run(['umount', disk],
                      stdout=devnull, stderr=devnull, as_root=True).wait()
            except cmd_runner.SubcommandNonZeroReturnValue:
                pass
    # Remove TMP_DIR as root because some files written there are
    # owned by root.
    if TMP_DIR is not None:
        cmd_runner.run(['rm', '-rf', TMP_DIR], as_root=True).wait()


def ensure_required_commands(args):
    """Ensure we have the commands that we know are going to be used."""
    required_commands = [
        'mkfs.vfat', 'sfdisk', 'mkimage', 'parted']
    for command in required_commands:
        ensure_command(command)


if __name__ == '__main__':
    parser = get_android_args_parser()
    args = parser.parse_args()

    logger = get_logger(debug=args.debug)

    additional_android_option_checks(args)

    # If --help was specified this won't execute.
    # Create temp dir and initialize rest of path vars.
    TMP_DIR = tempfile.mkdtemp()
    BOOT_DIR = os.path.join(TMP_DIR, 'boot')
    SYSTEM_DIR = os.path.join(TMP_DIR, 'system')
    DATA_DIR = os.path.join(TMP_DIR, 'data')

    BOOT_DISK = os.path.join(TMP_DIR, 'boot-disc')
    SYSTEM_DISK = os.path.join(TMP_DIR, 'system-disc')
    CACHE_DISK = os.path.join(TMP_DIR, 'cache-disc')
    DATA_DISK = os.path.join(TMP_DIR, 'userdata-disc')
    SDCARD_DISK = os.path.join(TMP_DIR, 'sdcard-disc')

    board_config = get_board_config(args.dev)
    if args.hwpack:
        board_config.from_file(args.hwpack)
    else:
        logger.warning("No Android hwpack provided: default board values "
                       "will be used.")
    board_config.add_boot_args(args.extra_boot_args)
    board_config.add_boot_args_from_file(args.extra_boot_args_file)

    if args.dev == 'iMX53':
        # XXX: remove this and the corresponding entry in android_board_configs
        logger.warning("DEPRECATION WARNING: iMX53 is deprecated, please "
                       "use mx53loco.")

    ensure_required_commands(args)

    media = Media(args.device)
    if media.is_block_device:
        if not confirm_device_selection_and_ensure_it_is_ready(args.device):
            sys.exit(1)
    elif not args.should_create_partitions:
        logger.error("Do not use --no-part in conjunction with --image_file.")
        sys.exit(1)

    cmd_runner.run(['mkdir', '-p', BOOT_DIR]).wait()
    cmd_runner.run(['mkdir', '-p', SYSTEM_DIR]).wait()
    cmd_runner.run(['mkdir', '-p', DATA_DIR]).wait()

    unpack_android_binary_tarball(args.boot, BOOT_DIR)
    unpack_android_binary_tarball(args.system, SYSTEM_DIR)
    unpack_android_binary_tarball(args.userdata, DATA_DIR)

    # Create partitions
    boot_partition, system_partition, cache_partition, \
        data_partition, sdcard_partition = setup_android_partitions( \
        board_config, media, args.image_size, args.boot_label,
        args.should_create_partitions, args.should_align_boot_part)

    board_config.populate_raw_partition(args.device, BOOT_DIR)
    populate_partition(BOOT_DIR + "/boot", BOOT_DISK, boot_partition)
    board_config.populate_boot_script(boot_partition, BOOT_DISK, args.consoles)
    with partition_mounted(boot_partition, BOOT_DISK):
        board_config.install_boot_loader(args.device, BOOT_DISK)
    populate_partition(SYSTEM_DIR + "/system", SYSTEM_DISK, system_partition)
    populate_partition(DATA_DIR + "/data", DATA_DISK, data_partition)
    print "Done creating Linaro Android image on %s" % args.device