#!/usr/bin/env python

# live-wrapper - Wrapper for vmdebootstrap for creating live images
# (C) Iain R. Learmonth 2015 <irl@debian.org>
# See COPYING for terms of usage, modification and redistribution.
#
# bin/lwr - Live Wrapper (Application)

"""
This script is the main script for the live-wrapper application. It is
intended to be run from the command line.

See live-wrapper(8) for more information.
"""

# pylint: disable=wrong-import-order,missing-docstring,superfluous-parens
# pylint: disable=too-many-statements,invalid-name

import sys
import os
import cliapp
import logging
import pycurl
import tempfile
import subprocess
from shutil import rmtree
from tarfile import TarFile
from lwr.vm import VMDebootstrap
from lwr.isolinux import install_isolinux, update_isolinux
from lwr.disk import install_disk_info
from lwr.grub import install_grub, update_grub
from lwr.xorriso import Xorriso
from lwr.apt_udeb import AptUdebDownloader
from lwr.utils import cdrom_image_url, KERNEL, RAMDISK

__version__ = '0.4'

# pylint: disable=line-too-long,too-many-function-args

# FIXME: make these available as lookups - single place to change.
directories = {
    'cdroot': '/',
    'boot': {
        'boot': '%s/boot/' % 'cdroot',  # fetch_di_helpers
        'boot/grub': '%s/boot/grub/' % 'cdroot',  # created by fetch_di_helpers
    },
    'isolinux': '%s/isolinux/' % 'cdroot',  # isolinux directory
    'd-i': {
        '%s/d-i/gtk' % 'cdroot',  # installer kernel & ramdisk
    },
    'pool': '%s/pool/' % 'cdroot',
    'dists': '%s/dists/' % 'cdroot',
    'metadata': '%s/.disk' % 'cdroot',  # disk.py
    'live': '%s/live/' % 'cdroot',
}

# FIXME: ISOLINUX looks in boot/isolinux, isolinux and then / - not just boot.
# splash.png needs to be in isolinux/

# Need to edit the D-I menu to add the Live! boot whilst allowing
# access to the D-I boot options.
# grub still not finding modules.


class LiveWrapper(cliapp.Application):

    # Instance variables
    cdroot = None  # The path to the chroot the CD is being built in
    kernel_path = None
    ramdisk_path = None
    gtk_kernel_path = None
    gtk_ramdisk_path = None

    def add_settings(self):
        default_arch = subprocess.check_output(["dpkg", "--print-architecture"]).decode('utf-8').strip()
        self.settings.string(
            ['o', 'image_output'], 'Location for built image',
            metavar='/PATH/TO/FILE.ISO',
            default='live.iso', group='Base Settings')
        self.settings.string(
            ['d', 'distribution'], 'Debian release to use (default: %default)',
            metavar='NAME',
            default='stretch', group='Base Settings')
        self.settings.string(
            ['architecture'], 'architecture to use (%default)',
            metavar='ARCH', default=default_arch)
        self.settings.string(
            ['m', 'mirror'], 'Mirror to use for image creation (default: %default)',
            metavar='MIRROR',
            group='Base Settings',
            default='http://ftp.debian.org/debian/')
        self.settings.string(
            ['t', 'tasks'], 'Task packages to install',
            metavar='"task-TASK1 task-TASK2 ..."',
            group='Packages')
        self.settings.string(
            ['e', 'extra'], 'Extra packages to install',
            metavar='"PKG1 PKG2 ..."',
            group='Packages')
        self.settings.boolean(
            ['isolinux'], 'Add isolinux bootloader to the image '
            '(default: %default)', default=True, group="Bootloaders")
        self.settings.boolean(
            ['grub'], 'Add GRUB bootloader to the image (for EFI support) '
            '(default: %default)', default=True, group="Bootloaders")
        self.settings.boolean(
            ['grub-loopback-only'], 'Only install the loopback.cfg GRUB '
            'configuration (for loopback support) (overrides --grub) '
            '(default: %default)', default=False, group="Bootloaders")
        self.settings.boolean(
            ['installer'], 'Include Debian Installer in the Live image',
            default=False, group="Debian Installer")

    def process_args(self, args):
        if os.path.exists(self.settings['image_output']):
            raise cliapp.AppException("Image '%s' already exists" % self.settings['image_output'])
        if not self.settings['isolinux'] and not self.settings['grub']:
            raise cliapp.AppException("You must enable at least one bootloader!")
        if self.settings['grub'] and self.settings['grub-loopback-only']:
            self.settings['grub'] = False
        if os.geteuid() != 0:
            sys.exit("You need to have root privileges to run this script.")
        # FIXME: cleanup on error.
        self.start_ops()

    def fetch_di_helpers(self, mirror, suite, architecture):
        logging.info("Downloading helper files from debian-installer team...")
        urls = cdrom_image_url(mirror, suite, architecture, gtk=False)
        bootdir = os.path.join(self.cdroot, 'boot')
        try:
            if self.settings['installer']:
                os.makedirs(os.path.join(self.cdroot, 'd-i', 'gtk'))
                self.kernel_path = os.path.join(self.cdroot, 'd-i', KERNEL)
                self.ramdisk_path = os.path.join(self.cdroot, 'd-i', RAMDISK)
            # python3 only
            if sys.version_info > (3, 0):
                ditar = tempfile.mkstemp()
                with ditar[0] as info:
                    curl = pycurl.Curl()
                    curl.setopt(curl.URL, urls[3])
                    curl.setopt(curl.WRITEDATA, info)
                    curl.perform()
                    curl.close()
                info = TarFile.open(ditar[1], 'r:gz')
                info.extractall(path=bootdir)
                info.close()
                os.remove(ditar[1])
            # python 2
            else:
                ditar = tempfile.NamedTemporaryFile(delete=False)  # pylint: disable=redefined-variable-type
                curl = pycurl.Curl()
                curl.setopt(curl.URL, urls[3])
                curl.setopt(curl.WRITEDATA, ditar)
                curl.perform()
                curl.close()
                ditar.close()
                info = TarFile.open(ditar.name, 'r:gz')
                info.extractall(path=bootdir)
                info.close()
                os.remove(ditar.name)

            # fetch debian-installer
            if self.settings['installer']:
                with open(self.kernel_path, 'w') as kernel:
                    curl = pycurl.Curl()
                    curl.setopt(curl.URL, urls[1])
                    curl.setopt(curl.WRITEDATA, kernel)
                    curl.perform()
                    curl.close()
                with open(self.ramdisk_path, 'w') as ramdisk:
                    curl = pycurl.Curl()
                    curl.setopt(curl.URL, urls[2])
                    curl.setopt(curl.WRITEDATA, ramdisk)
                    curl.perform()
                    curl.close()
        except pycurl.error:
            logging.error("Failed to fetch the debian-installer helper files! " +
                          "Cannot continue!")
            sys.exit(1)
        # Now get the graphical installer.
        if self.settings['installer']:
            urls = cdrom_image_url(mirror, suite, architecture, gtk=True)
            self.gtk_kernel_path = os.path.join(self.cdroot, 'd-i', 'gtk', KERNEL)
            self.gtk_ramdisk_path = os.path.join(self.cdroot, 'd-i', 'gtk', RAMDISK)
            try:
                with open(self.gtk_kernel_path, 'w') as kernel:
                    curl = pycurl.Curl()
                    curl.setopt(curl.URL, urls[1])
                    curl.setopt(curl.WRITEDATA, kernel)
                    curl.perform()
                    curl.close()
                with open(self.gtk_ramdisk_path, 'w') as ramdisk:
                    curl = pycurl.Curl()
                    curl.setopt(curl.URL, urls[2])
                    curl.setopt(curl.WRITEDATA, ramdisk)
                    curl.perform()
                    curl.close()
            except pycurl.error:
                logging.error("Failed to fetch the gtk debian-installer helper files! " +
                              "Cannot continue!")
                sys.exit(1)

    def start_ops(self):  # pylint: disable=too-many-statements
        """
        This function creates the live image using the settings determined by
        the arguments passed on the command line.

        .. note::
            This function is called by process_args() once all the arguments
            have been validated.
        """

        print("Creating work directory...")
        logging.info("Creating work directory...")

        # Create work directory
        self.cdroot = tempfile.mkdtemp()  # all other directories are based off this

        logging.debug("Setting environment variables for customise hook...")

        # Make options available to customise hook in vmdebootstrap
        os.environ['LWR_MIRROR'] = self.settings['mirror']
        os.environ['LWR_DISTRIBUTION'] = self.settings['distribution']
        os.environ['LWR_TASK_PACKAGES'] = self.settings['tasks']
        os.environ['LWR_EXTRA_PACKAGES'] = self.settings['extra']

        logging.info("Running vmdebootstrap...")
        print("Running vmdebootstrap...")

        # Run vmdebootstrap, putting files in /live/
        vm = VMDebootstrap(self.settings['distribution'],
                           self.settings['architecture'],
                           self.settings['mirror'], self.cdroot)
        vm.run()

        # Fetch D-I helper archive if needed
        if self.settings['grub']:
            print("Fetching Debian Installer helpers")
            self.fetch_di_helpers(
                self.settings['mirror'],
                self.settings['distribution'],
                self.settings['architecture'])

        # Download the udebs
        if self.settings['installer']:
            print("Downloading udebs for Debian Installer...")  # FIXME: self.message()
            logging.info("Downloading udebs for Debian Installer...")
            # FIXME: get exclude_list from user
            exclude_list = []
            # FIXME: may need a change to the download location
            di_root = os.path.join(self.cdroot, 'd-i')
            apt_udeb = AptUdebDownloader(destdir=di_root)
            apt_udeb.mirror = self.settings['mirror']
            apt_udeb.architecture = self.settings['architecture']
            apt_udeb.suite = self.settings['distribution']
            print("Updating a local cache for %s %s ..." % (apt_udeb.architecture, apt_udeb.suite))  # FIXME: self.message()
            logging.debug("Updating local cache...")
            apt_udeb.prepare_apt()
            # FIXME: add support for a custom apt source on top.
    
            # download all udebs in the suite, except exclude_list
            apt_udeb.download_udebs(exclude_list)
            # FIXME: generate a Release file
            apt_udeb.clean_up_apt()
            print("... completed udeb downloads")
            logging.info("... completed udeb downloads")

        # Install isolinux if selected
        if self.settings['isolinux']:
            logging.info("Performing isolinux installation...")
            bootdir = os.path.join(self.cdroot, 'isolinux')
            os.mkdir(bootdir)
            # FIXME: catch errors and cleanup.
            install_isolinux(
                bootdir, self.settings['mirror'],
                self.settings['distribution'],
                self.settings['architecture'])
            if self.settings['installer']:
                update_isolinux(bootdir, self.kernel_path, self.ramdisk_path)

        # Install GRUB if selected
        if self.settings['grub'] or self.settings['grub-loopback-only']:
            logging.info("Performing GRUB installation...")
            install_grub(self.cdroot) # FIXME: pass architecture & uefi settings.
            if self.settings['installer']:
                update_grub(self.cdroot, self.kernel_path, self.ramdisk_path)

        # Install .disk information
        logging.info("Installing the disk metadata ...")
        install_disk_info(self.cdroot)

        # Create ISO image
        logging.info("Creating the ISO image with Xorriso...")
        xorriso = Xorriso(self.settings['image_output'],
                          isolinux=self.settings['isolinux'],
                          grub=self.settings['grub'])
        xorriso.build_args(self.cdroot)
        xorriso.build_image()

        # Remove the temporary directories
        logging.info("Removing temporary work directories...")
        if self.settings['installer']:
            apt_udeb.clean_up_apt()
        # rmtree(self.cdroot)
        print("Use the -cdrom option to test the image using qemu-system.")

if __name__ == "__main__":
    LiveWrapper(version=__version__).run()
