#!/usr/bin/env python

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

from linaro_media_create.boards import (
    board_configs,
    ROOTFS_UUID,
    )
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,
    )
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


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')


# 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.
    Popen(
        ['sudo', 'umount', BOOT_DISK], stdout=devnull, stderr=devnull).wait()
    Popen(
        ['sudo', 'umount', ROOT_DISK], stdout=devnull, stderr=devnull).wait()
    # Remove TMP_DIR as root because some files written there are
    # owned by root.
    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', 'fdisk', 'mkimage', 'uuidgen', '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()
    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(
        args.board, media, board_config.fat_size, args.image_size,
        args.boot_label, args.rfs_label, args.rootfs, ROOTFS_UUID,
        args.should_create_partitions,
        args.should_format_bootfs, args.should_format_rootfs)

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