#!/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 board_configs
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.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 ensure_command, is_arm_host

# 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.
@atexit.register
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')
        required_commands.append('qemu-img')
    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:
        ensure_command(command)


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

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

    board_config = board_configs[args.board]

    ensure_required_commands(args)

    sig_file_list = args.hwpacksigs[:]
    if args.binarysig is not None:
        sig_file_list.append(args.binarysig)
    for sig_file in sig_file_list:
        hash_file = sig_file[0:-len('.asc')]
        if cmd_runner.run(['gpg', '--verify', sig_file]).wait() != 0:
            print "Could not verify hash file signature %s." % sig_file
            sys.exit(1)
        if cmd_runner.run(['sha1sum', '-c', hash_file]).wait() != 0:
            print "Could not verify hash in file %s." % hash_file
            sys.exit(1)

    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_format_rootfs or not args.should_format_bootfs:
        print ("Do not use --no-boot or --no-part in conjunction with "
               "--image_file.")
        sys.exit(1)
    else:
        # All good, move on.
        pass

    unpack_binary_tarball(args.binary, TMP_DIR)

    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, *hwpacks)

    if args.rootfs == 'btrfs':
        install_packages(ROOTFS_DIR, TMP_DIR, "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)

    rootfs_uuid = get_uuid(root_partition)

    if args.should_format_bootfs:
        board_config.populate_boot(
            ROOTFS_DIR, rootfs_uuid, boot_partition, BOOT_DISK, args.device,
            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_uuid, create_swap, str(args.swap_file),
            board_config.mmc_part_offset)

    print "Done creating Linaro image on %s" % args.device
