#!/usr/bin/env python3
# This file is part of Cockpit.
#
# Copyright (C) 2015 Red Hat, Inc.
#
# Cockpit is free software; you can redistribute it and/or modify it
# under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation; either version 2.1 of the License, or
# (at your option) any later version.
#
# Cockpit 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
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with Cockpit; If not, see <http://www.gnu.org/licenses/>.

import argparse
import errno
import os
import sys
import subprocess

import make_dist

from common.parent import BASE_DIR, TEST_DIR, BOTS_DIR, ensure_bots
ensure_bots()  # NOQA: testvm lives in bots/
import testvm

os.environ["PATH"] = "{0}:{1}".format(os.environ.get("PATH"), BOTS_DIR)

networking = None


def main():
    parser = argparse.ArgumentParser(
        description='Prepare testing environment, download images and build and install cockpit',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('-v', '--verbose', action='store_true', help='Display verbose progress details')
    parser.add_argument('-s', '--sit', action='store_true', help='Sit and wait if install script fails')
    parser.add_argument('-q', '--quick', action='store_true', help='Skip unit tests to build faster')
    parser.add_argument('-o', '--overlay', action='store_true', help='Install into existing test/image/ overlay instead of from pristine base image')
    parser.add_argument('-b', '--build-image', action='store', help='Build in this image')
    parser.add_argument('-B', '--build-only', action='store_true', help='Only build and download results')
    parser.add_argument('-I', '--install-only', action='store_true', help='Only upload and install')
    parser.add_argument('-c', '--containers', action='store_true', help='Install container images')
    parser.add_argument('--address', help='Address of already running machine, implies --install-only')
    parser.add_argument('image', nargs='?', default=testvm.DEFAULT_IMAGE, help='The image to use')
    args = parser.parse_args()

    # The basic operating system we're preparing on top of. Never written to
    base_image = args.image

    # The image we're going to build in. This image is never written to
    build_image = None
    if not args.install_only:
        build_image = testvm.get_build_image(args.build_image or base_image)

    # Depending on the operating system ignore some things
    skips = []
    if args.address:
        skips.append("cockpit-tests")
        args.install_only = True

    if args.install_only and args.build_only:
        sys.stderr.write("image-prepare: use of --install-only with --build-only is incompatible\n")
        return 2

    results = os.path.abspath("tmp/build-results")

    # Make sure any images are downloaded
    try:
        if not args.address and not os.path.exists(base_image):
            subprocess.check_call(["image-download", base_image])
        if build_image and not os.path.exists(build_image):
            subprocess.check_call(["image-download", build_image])
    except OSError as ex:
        if ex.errno != errno.ENOENT:
            raise
        sys.stderr.write("image-prepare: missing tools to download images\n")
    except subprocess.CalledProcessError:
        sys.stderr.write("image-prepare: unable to download all necessary images\n")
        return 1

    try:
        if not args.install_only:
            build_and_install(build_image, results, skips, args)
        if not args.build_only and build_image != base_image:
            only_install(base_image, results, skips, args)
    except (RuntimeError, testvm.Failure) as ex:
        sys.stderr.write("image-prepare: {0}\n".format(str(ex)))
        return 2

    return 0


def upload_scripts(machine, args):
    machine.execute("rm -rf /var/lib/testvm")
    machine.upload([os.path.join(testvm.SCRIPTS_DIR, "lib")], "/var/lib/testvm")
    machine.upload([os.path.join(BASE_DIR, "tools", "make-srpm")], "/var/lib/testvm")
    machine.upload([os.path.join(testvm.SCRIPTS_DIR, "%s.install" % machine.image)], "/var/tmp")
    machine.upload([os.path.join(BASE_DIR, "containers")], "/var/tmp")


# Create the necessary layered image for the build/install
def prepare_install_image(base_image, overlay=False):
    if "/" not in base_image:
        base_image = os.path.join(testvm.IMAGES_DIR, base_image)
    test_images = os.path.join(TEST_DIR, "images")
    os.makedirs(test_images, exist_ok=True)
    install_image = os.path.join(test_images, os.path.basename(base_image))
    if not os.path.exists(install_image) or not overlay:
        base_image = os.path.realpath(testvm.get_test_image(base_image))
        qcow2_image = "{0}.qcow2".format(install_image)
        subprocess.check_call(["qemu-img", "create", "-q", "-f", "qcow2",
                               "-o", "backing_file={0},backing_fmt=qcow2".format(base_image),
                               qcow2_image])
        return install_image, qcow2_image
    else:
        return install_image, None


def run_install_script(machine, do_build, do_install, skips, arg, args):
    install = do_install
    if args.containers and not do_build:
        do_install = False

    skips = list(skips or [])

    skip_args = [" --skip '%s'" % skip for skip in skips]
    cmd = "cd /var/tmp; ./%s.install%s%s%s%s%s%s" % (machine.image,
                                                     " --verbose" if args.verbose else "",
                                                     " --quick" if args.quick else "",
                                                     " --build" if do_build else "",
                                                     " --install" if do_install else "",
                                                     " ".join(skip_args),
                                                     " '%s'" % arg if arg else "")
    machine.execute(cmd, timeout=1800)

    if install and args.containers:
        machine.execute("/var/lib/testvm/containers.install", timeout=900)


def finalize_image(target, qcow_overlay):
    '''Create final file name

    If the target image does not exist already, prepare_install_image() creates
    a <target>.qcow2 overlay. Symlink this to the final name (without .qcow2) at the
    end to avoid leaving half-broken images around in the case of errors.
    '''
    if qcow_overlay:
        if os.path.lexists(target):
            os.unlink(target)
        os.symlink(os.path.basename(qcow_overlay), target)


def build_and_install(build_image, build_results, skips, args):
    """Build and maybe install Cockpit into a test image"""
    target, overlay = prepare_install_image(build_image, args.overlay)
    machine = testvm.VirtMachine(verbose=args.verbose, image=(overlay or target),
                                 memory_mb=2048, cpus=4, maintain=True, networking=networking)
    completed = False

    # While the machine is booting, make a dist tarball
    source = make_dist.make_dist()

    machine.start()
    completed = False

    try:
        machine.wait_boot()
        upload_scripts(machine, args=args)
        machine.upload([source], "/var/tmp")
        run_install_script(machine, True, True, skips, os.path.basename(source), args)
        finalize_image(target, overlay)
        completed = True
    finally:
        if not completed and args.sit:
            sys.stderr.write(machine.diagnose())
            input("Press RET to continue... ")
        try:
            if os.path.exists(build_results):
                subprocess.check_call(["rm", "-rf", build_results])
            os.makedirs(build_results)
            machine.download("/var/tmp/build-results/*", build_results)
        finally:
            machine.stop()


def only_install(image, build_results, skips, args):
    """Install Cockpit into a test image"""
    started = False
    if args.address:
        machine = testvm.Machine(address=args.address, verbose=args.verbose, image=image)
    else:
        target, overlay = prepare_install_image(image, args.overlay)
        machine = testvm.VirtMachine(verbose=args.verbose, image=(overlay or target), networking=networking, maintain=True)
        machine.start()
        started = True
    completed = False
    try:
        if started:
            machine.wait_boot()
        upload_scripts(machine, args=args)
        machine.execute("rm -rf /var/tmp/build-results")
        machine.upload([build_results], "/var/tmp/build-results")
        run_install_script(machine, False, True, skips, None, args)
        finalize_image(target, overlay)
        completed = True
    finally:
        if not completed and args.sit:
            sys.stderr.write(machine.diagnose())
            input("Press RET to continue... ")
        if started:
            machine.stop()


if __name__ == "__main__":
    sys.exit(main())
