#!/usr/bin/env python
# Copyright (C) 2010, 2011 Linaro
#
# Author: Guilherme Salgado <guilherme.salgado@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.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.chroot_utils import (
    install_hwpacks,
    install_packages,
    )
from linaro_image_tools.hwpack.hwpack_reader import (
    HwpackReader,
    HwpackReaderError,
    )
from linaro_image_tools.media_create.partitions import (
    Media,
    setup_partitions,
    get_uuid,
    )
from linaro_image_tools.media_create.rootfs import populate_rootfs
from linaro_image_tools.media_create.unpack_binary_tarball import (
    unpack_binary_tarball,
    )
from linaro_image_tools.media_create import get_args_parser
from linaro_image_tools.utils import (
    additional_option_checks,
    check_file_integrity_and_log_errors,
    check_required_args,
    ensure_command,
    IncompatibleOptions,
    is_arm_host,
    MissingRequiredOption,
    path_in_tarfile_exists,
    prep_media_path,
    get_logger,
    UnableToFindPackageProvidingCommand,
    disable_automount,
    enable_automount,
    )

# Just define the global variables
TMP_DIR = None
ROOTFS_DIR = None
BOOT_DISK = None
ROOT_DISK = None


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

    Before doing so, make sure BOOT_DISK and ROOT_DISK are not mounted.
    """
    devnull = open('/dev/null', 'w')
    # ignore non-zero return codes
    for disk in BOOT_DISK, ROOT_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', 'gpg', 'sha1sum']
    if not is_arm_host():
        required_commands.append('qemu-arm-static')
    if args.rootfs in ['btrfs', 'ext2', 'ext3', 'ext4']:
        required_commands.append('mkfs.%s' % args.rootfs)
    else:
        raise AssertionError('Unsupported rootfs type %s' % args.rootfs)

    for command in required_commands:
        try:
            ensure_command(command)
        except UnableToFindPackageProvidingCommand:
            logger.error("Could not look up command %s. Please ensure that command %s is installed." % (command, command))
            raise


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

    logger = get_logger(debug=args.debug)

    try:
        additional_option_checks(args)
    except IncompatibleOptions as e:
        parser.print_help()
        logger.error(e.value)
        sys.exit(1)

    if args.readhwpack:
        try:
            reader = HwpackReader(args.hwpacks)
            logger.info(reader.get_supported_boards())
            sys.exit(0)
        except HwpackReaderError as e:
            logger.error(e.value)
            sys.exit(1)

    try:
        check_required_args(args)
    except MissingRequiredOption as e:
        parser.print_help()
        logger.error(e.value)
        sys.exit(1)

    # Do this by default, disable automount options and re-enable them at exit.
    disable_automount()
    atexit.register(enable_automount)

    board_config = get_board_config(args.dev)
    board_config.set_metadata(args.hwpacks, args.bootloader, args.dev)
    board_config.add_boot_args(args.extra_boot_args)
    board_config.add_boot_args_from_file(args.extra_boot_args_file)

    media = Media(prep_media_path(args))

    if media.is_block_device:
        if not board_config.supports_writing_to_mmc:
            logger.error("The board '%s' does not support the --mmc option. "
                         "Please use --image_file to create an image file for "
                         "this board." % args.dev)
            sys.exit(1)
        if not confirm_device_selection_and_ensure_it_is_ready(
                args.device, args.nocheck_mmc):
            sys.exit(1)
    elif not args.should_format_rootfs or not args.should_format_bootfs:
        logger.error("Do not use --no-boot or --no-part in conjunction with "
                     "--image_file.")
        sys.exit(1)

    # If --help was specified this won't execute.
    # Create temp dir and initialize rest of path vars.
    TMP_DIR = tempfile.mkdtemp()
    BOOT_DISK = os.path.join(TMP_DIR, 'boot-disc')
    ROOT_DISK = os.path.join(TMP_DIR, 'root-disc')
    BIN_DIR = os.path.join(TMP_DIR, 'rootfs')
    os.mkdir(BIN_DIR)

    logger.info('Searching correct rootfs path')
    # Identify the correct path for the rootfs
    filesystem_dir = ''
    if path_in_tarfile_exists('binary/etc', args.binary):
        filesystem_dir = 'binary'
    elif path_in_tarfile_exists('binary/boot/filesystem.dir', args.binary):
        # The binary image is in the new live format.
        filesystem_dir = 'binary/boot/filesystem.dir'

    ROOTFS_DIR = os.path.join(BIN_DIR, filesystem_dir)

    try:
        ensure_required_commands(args)
    except UnableToFindPackageProvidingCommand:
        sys.exit(1)

    sig_file_list = args.hwpacksigs[:]
    if args.binarysig is not None:
        sig_file_list.append(args.binarysig)

    # Check that the signatures that we have been provided (if any) match
    # the hwpack and OS binaries we have been provided. If they don't, quit.
    files_ok, verified_files = check_file_integrity_and_log_errors(
                                    sig_file_list, args.binary, args.hwpacks)
    if not files_ok:
        sys.exit(1)

    atexit.register(cleanup_tempdir)

    unpack_binary_tarball(args.binary, BIN_DIR)

    # if compatible system, extract all packages
    os_release_id = 'linux'
    os_release_file = '%s/etc/os-release' % ROOTFS_DIR
    if os.path.exists(os_release_file):
        for line in open(os_release_file):
            if line.startswith('ID='):
                os_release_id = line[(len('ID=')):]
                os_release_id = os_release_id.strip('\"\n')
                break

    if os_release_id == 'debian' or os_release_id == 'ubuntu' or \
            os.path.exists('%s/etc/debian_version' % ROOTFS_DIR):
        extract_kpkgs = False
    elif os_release_id == 'fedora':
        extract_kpkgs = False
    else:
        extract_kpkgs = True

    hwpacks = args.hwpacks
    lmc_dir = os.path.dirname(__file__)
    if lmc_dir == '':
        lmc_dir = None
    install_hwpacks(ROOTFS_DIR, TMP_DIR, lmc_dir, args.hwpack_force_yes,
                    verified_files, extract_kpkgs, *hwpacks)

    if args.rootfs == 'btrfs':
        if not extract_kpkgs:
            logger.info("Desired rootfs type is 'btrfs', trying to "
                        "auto-install the 'btrfs-tools' package")
            install_packages(ROOTFS_DIR, TMP_DIR, "btrfs-tools")
        else:
            logger.info("Desired rootfs type is 'btrfs', please make sure the "
                        "rootfs also includes 'btrfs-tools'")

    boot_partition, root_partition = setup_partitions(
        board_config, media, args.image_size, args.boot_label, args.rfs_label,
        args.rootfs, args.should_create_partitions, args.should_format_bootfs,
        args.should_format_rootfs, args.should_align_boot_part)

    uuid = get_uuid(root_partition)
    # In case we're only extracting the kernel packages, avoid
    # using uuid because we don't have a working initrd
    if extract_kpkgs:
        # XXX: this needs to be smarter as we can't always assume mmcblk devices
        rootfs_id = '/dev/mmcblk%dp%s' % (
                board_config.mmc_device_id, 2 + board_config.mmc_part_offset)
    else:
        rootfs_id = "UUID=%s" % uuid

    if args.should_format_bootfs:
        board_config.populate_boot(
            ROOTFS_DIR, rootfs_id, boot_partition, BOOT_DISK, media.path,
            args.is_live, args.is_lowmem, args.consoles)

    if args.should_format_rootfs:
        create_swap = False
        if args.swap_file is not None:
            create_swap = True
        populate_rootfs(ROOTFS_DIR, ROOT_DISK, root_partition, args.rootfs,
            rootfs_id, create_swap, str(args.swap_file),
            board_config.mmc_device_id, board_config.mmc_part_offset,
            os_release_id, board_config)

    logger.info("Done creating Linaro image on %s" % media.path)