#!/usr/bin/python3
# -*- coding: utf-8 -*-

# Copyright (C) 2013 Canonical Ltd.
# Author: Stéphane Graber <stgraber@ubuntu.com>

# This program 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; version 3 of the License.
#
# This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.

import argparse
import glob
import lxc
import os
import stat
import shutil
import subprocess
import sys
import time
import platform

from distro_info import UbuntuDistroInfo

CONTAINER_NAME = "unity8-lxc"
ISO_URL = "http://cdimage.ubuntu.com/ubuntu-desktop-next/daily-live/" \
          "current/RELEASE-desktop-ARCH.iso.zsync"
SKIP_JOB = ["bluetooth", "lightdm", "plymouth", "plymouth-upstart-bridge",
            "plymouth-shutdown", "ubiquity"]


# Argument parsing
parser = argparse.ArgumentParser(description="Unity8 inside LXC")
parser.add_argument(
    "--rebuild-all", action="store_true",
    help=("Wipe and replace the container rootfs and configuration."))
parser.add_argument(
    "--rebuild-config", action="store_true",
    help=("Wipe and regenerate the container configuration."))
parser.add_argument(
    "--rebuild-rootfs", action="store_true",
    help=("Wipe and replace the container rootfs."))
parser.add_argument(
    "--redownload", action="store_true",
    help=("Re-download the ISO image instead of using the local version."))
parser.add_argument(
    "--test", type=int, metavar="SECONDS", default=None,
    help=("Test mode, argument is time in seconds for the session."))
parser.add_argument(
    "--destroy", action="store_true",
    help=("Destroy any existing environment entirely."))
parser.add_argument(
    "--update-lxc", action="store_true",
    help=("Update the packages in the Unity 8 desktop preview LXC."))
args = parser.parse_args()

if os.geteuid():
    parser.error("This tool must be run as root.")


container = lxc.Container(CONTAINER_NAME)
container_path = "%s/%s" % (lxc.default_config_path, CONTAINER_NAME)
rootfs_path = "%s/rootfs" % container_path
iso_filename = "%s/ubuntu-next.iso" % container_path
iso_path = "%s/iso" % container_path
squashfs_path = "%s/squashfs" % container_path

(distro, version, codename) = platform.linux_distribution()

# Actions
generate_config = True
generate_rootfs = True

# Destroy a container entirely
if args.destroy and container.defined:
    container.stop()
    container.destroy()
    sys.exit(0)

# Deal with existing containers
if container.defined:
    generate_config = False
    generate_rootfs = False

    if not args.test and container.running:
        print("Stopping the existing container.")
        container.stop()

    if args.rebuild_all:
        args.rebuild_config = True
        args.rebuild_rootfs = True

    if args.rebuild_config:
        generate_config = True
        container.clear_config()

    if args.rebuild_rootfs:
        generate_rootfs = True
        if os.path.exists(rootfs_path):
            shutil.rmtree(rootfs_path)

if not generate_config and not generate_rootfs and not args.test and not args.update_lxc:
    parser.error("The container already exists.")


if generate_config:
    # Setup the container config
    ## Load the default LXC config keys
    container.load_config("/usr/share/lxc/config/ubuntu.common.conf")

    ## Setup the network ##
    container.append_config_item("lxc.network.type", "veth")
    container.append_config_item("lxc.network.link", "lxcbr0")
    container.append_config_item("lxc.network.hwaddr", "00:16:3e:xx:xx:xx")

    ## Architecture
    container.set_config_item("lxc.arch", "x86_64")

    ## Rootfs path
    container.set_config_item("lxc.rootfs", rootfs_path)

    ## Hostname
    container.set_config_item("lxc.utsname", CONTAINER_NAME)

    ## Devices
    ### /dev/tty*
    container.append_config_item("lxc.cgroup.devices.allow", "c 4:* rwm")

    ### /dev/input/*
    container.append_config_item("lxc.cgroup.devices.allow", "c 13:* rwm")
    container.append_config_item("lxc.mount.entry",
                                 "/dev/input dev/input none bind,create=dir")

    ### /dev/dri/*
    container.append_config_item("lxc.cgroup.devices.allow", "c 226:* rwm")
    container.append_config_item("lxc.mount.entry",
                                 "/dev/dri dev/dri none bind,create=dir")

    ### /dev/snd/*
    container.append_config_item("lxc.cgroup.devices.allow", "c 116:* rwm")
    container.append_config_item("lxc.mount.entry",
                                 "/dev/snd dev/snd none bind,create=dir")

    ## Enable cgmanager access
    container.set_config_item("lxc.mount.auto", "cgroup:mixed")

    ## Automount /sys in 15.04 and later
    if version >= '15.04':
        container.append_config_item("lxc.mount.auto", "sys")

    ## Allow nested containers
    container.set_config_item("lxc.aa_profile", "lxc-container-default-with-nesting")

    ## Disable the lxc autodev
    container.append_config_item("lxc.autodev", "0")

    # Setup the /home bind-mount
    container.append_config_item("lxc.mount.entry", "/home home none bind")

    # Setup the /etc/shadow bind-mount
    container.append_config_item("lxc.mount.entry",
                                 "/etc/shadow etc/shadow none bind")

    # Setup the local time and timezone
    container.append_config_item("lxc.mount.entry",
                                 "/etc/localtime etc/localtime none bind")
    container.append_config_item("lxc.mount.entry",
                                 "/etc/timezone etc/timezone none bind")

    # Setup up /run/udev
    container.append_config_item("lxc.mount.entry",
                                 "/run/udev var/lib/host-udev none bind,ro,create=dir")

    # Setup /run/systemd
    container.append_config_item("lxc.mount.entry",
                                 "/run/systemd var/lib/host-systemd none bind,ro,create=dir")

    ## Dump it all to disk
    container.save_config()


if generate_rootfs:
    # Setup the container rootfs

    ## Figure out the host architecture
    dpkg = subprocess.Popen(['dpkg', '--print-architecture'],
                            stdout=subprocess.PIPE,
                            universal_newlines=True)
    if dpkg.wait() != 0:
        parser.error("Failed to determine the local architecture.")

    architecture = dpkg.stdout.read().strip()

    ## Download the ISO image
    if args.redownload or not os.path.exists(iso_filename):
        ### Figure out the latest Ubuntu development release
        devel_release = UbuntuDistroInfo().devel()

        iso_url = ISO_URL.replace("ARCH", architecture)
        iso_url = iso_url.replace("RELEASE", devel_release)

        ### Use zsync to fetch the ISO
        cmd = "cd %s && zsync -i %s -o %s %s" % (container_path, iso_filename, iso_filename, iso_url)
        os.system(cmd)

        ### Remove the ISO backup since it's not needed
        if os.path.exists(iso_filename+".zs-old"):
            os.remove(iso_filename+".zs-old")

    if not os.path.exists(iso_filename):
        parser.error("No Unity 8 Desktop Next ISO exists!")

    ## Mount the ISO
    os.mkdir(iso_path)
    subprocess.call(["mount", "-o", "loop,ro", iso_filename, iso_path])

    os.mkdir(squashfs_path)
    subprocess.call(["mount", "-o", "loop,ro",
                     os.path.join(iso_path, "casper", "filesystem.squashfs"),
                     squashfs_path])

    ## Unpack the ISO
    print("Unpacking the ISO image...")
    subprocess.call(["rsync", "-aA", "--hard-links", "--numeric-ids",
                     "%s/" % squashfs_path, "%s/" % rootfs_path])

    ## Unmount the ISO
    subprocess.call(["umount", squashfs_path])
    os.rmdir(squashfs_path)

    subprocess.call(["umount", iso_path])
    os.rmdir(iso_path)

    ## Configure
    print("Configuring the Unity8 LXC...")

    ### Generate /etc/hostname
    with open(os.path.join(rootfs_path, "etc", "hostname"), "w+") as fd:
        fd.write("unity8-lxc\n")

    ### Generate /etc/hosts
    with open(os.path.join(rootfs_path, "etc", "hosts"), "w+") as fd:
        fd.write("""127.0.0.1   localhost
127.0.1.1   unity8-mir

# The following lines are desirable for IPv6 capable hosts
::1     localhost ip6-localhost ip6-loopback
ff02::1 ip6-allnodes
ff02::2 ip6-allrouters
""")

    ### Disable some upstart jobs
    for job in glob.glob("%s/etc/init/*.conf" % rootfs_path):
        if os.path.basename(job).rsplit(".conf", 1)[0] in SKIP_JOB:
            with open("%s.override" % job.rsplit(".conf", 1)[0], "w+") as fd:
                fd.write("manual")

    ### Create any missing devices
    for i in range(64):
        tty_path = os.path.join(rootfs_path, "dev", "tty%s" % i)
        if not os.path.exists(tty_path):
            os.mknod(tty_path, 0o620 | stat.S_IFCHR, device=os.makedev(4, i))
            os.chown(tty_path, 0, 5)

    ### Generate /run/udev
    with open(os.path.join(rootfs_path, "etc", "init", "udev-db.conf"), "w+") as fd:
        fd.write("""start on starting udev
pre-start script
    rm -Rf /run/udev
    cp -R /var/lib/host-udev /run/udev
end script
""")

if args.test:
    # Start a test session

    if not container.running:
        print("Starting the container")
        if not container.start():
            parser.error("Unable to start the container.")
        container.wait("RUNNING")

    # Retrieve the current VT number
    fgconsole = subprocess.Popen(['fgconsole'],
                                 stdout=subprocess.PIPE,
                                 universal_newlines=True)
    if fgconsole.wait() != 0:
        parser.error("Failed to get current VT number.")

    current_vt = fgconsole.stdout.read().strip()

    def start_unity8():
        ## Setup the environment
        os.environ['HOME'] = "/root"
        os.environ['QT_QPA_PLATFORM'] = "ubuntumirclient"
        os.environ['XDG_RUNTIME_DIR'] = "/run/user/0/"
        os.environ['DESKTOP_SESSION'] = "unity8-mir"
        os.environ['UBUNTU_PLATFORM_API_BACKEND'] = "desktop_mirserver"

        ## Reset the config
        for path in ("/root/.cache", "/root/.config", "/root/.local"):
            if os.path.exists(path):
                shutil.rmtree(path)

        ## Create runtime dir
        if not os.path.exists("/run/user/0"):
            os.makedirs("/run/user/0")

        ## Get the multiarch triplet
        gcc = subprocess.Popen(['gcc', '--print-multiarch'],
                               stdout=subprocess.PIPE,
                               universal_newlines=True)
        if gcc.wait() != 0:
            parser.error("Failed to determine the multiarch triplet.")

        triplet = gcc.stdout.read().strip()

        ## Run url-dispatcher
        dispatcher_path = "/usr/lib/%s/url-dispatcher/" \
                          "update-directory" % triplet

        subprocess.call([dispatcher_path, "/usr/share/url-dispatcher/urls/"])

        ## Start the session
        unity8 = subprocess.Popen(["openvt", "-s", "-f", "-w", "-c", "12",
                                   "--",
                                   "init", "--user"])
        time.sleep(args.test)
        unity8.kill()
        unity8.wait()

        ## Switch back to original VT
        subprocess.call(["chvt", current_vt])

    container.attach_wait(start_unity8)

    ## Stopping the container
    container.stop()

if args.update_lxc:
    # Update packages inside the LXC

    if not container.running:
        print("Starting the container")
        if not container.start():
            parser.error("Unable to start the container.")
        container.wait("RUNNING")

    if not container.get_ips(interface='eth0', timeout=30):
        print("Not able to connect to the network.")
        sys.exit(0)

    container.attach_wait(lxc.attach_run_command,
                          ["umount", "/etc/localtime"])

    container.attach_wait(lxc.attach_run_command,
                          ["umount", "/etc/timezone"])

    print("Updating packages inside the LXC...")

    container.attach_wait(lxc.attach_run_command,
                          ["apt-get", "update"])

    container.attach_wait(lxc.attach_run_command,
                          ["apt-get", "dist-upgrade", "-y"])

    ## Stopping the container
    container.stop()
