#!/usr/bin/env python

import atexit
import os
import sys
import tempfile
from subprocess import Popen

from linaro_media_create.boards import (
    board_configs,
    )
from linaro_media_create.check_device import (
    confirm_device_selection_and_ensure_it_is_ready)
from linaro_media_create.hwpack import install_hwpacks
from linaro_media_create.partitions import (
    Media,
    setup_partitions,
    get_uuid,
    )
from linaro_media_create.populate_boot import populate_boot
from linaro_media_create.rootfs import populate_rootfs
from linaro_media_create.unpack_binary_tarball import unpack_binary_tarball
from linaro_media_create.utils import (
    ensure_command,
    is_arm_host,
    )
from linaro_media_create import get_args_parser

# 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')
    # Use raw subprocess.Popen as we don't want to stop in case the
    # commands end with a non-zero return code.
    if BOOT_DISK is not None:
        Popen(['sudo', 'umount', BOOT_DISK],
              stdout=devnull, stderr=devnull).wait()
    if ROOT_DISK is not None:
        Popen(['sudo', 'umount', ROOT_DISK],
              stdout=devnull, stderr=devnull).wait()
    # Remove TMP_DIR as root because some files written there are
    # owned by root.
    if TMP_DIR is not None:
        Popen(['sudo', 'rm', '-rf', TMP_DIR]).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']
    if not is_arm_host():
        required_commands.append('qemu-arm-static')
        required_commands.append('qemu-img')
    if args.rootfs in ['ext2', 'ext3', 'ext4']:
        required_commands.append('mkfs.%s' % args.rootfs)
    else:
        required_commands.append('mkfs.btrfs')
    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)

    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
    install_hwpacks(ROOTFS_DIR, TMP_DIR, args.hwpack_force_yes, *hwpacks)

    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)

    rootfs_uuid = get_uuid(root_partition)

    if args.should_format_bootfs:
        populate_boot(
            board_config, 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