#!/usr/bin/python -u
# -*- coding: utf-8 -*-

# Copyright (C) 2013 Stéphane Graber
# 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/>.

from __future__ import division, print_function

import argparse
import fcntl
import os
import urllib
import subprocess
import shutil
import tempfile
import re

from ubuntutools.misc import host_architecture

from ConfigParser import ConfigParser

cache_dir = None


def guess_backingstore(path):
    with open(os.devnull, "w") as devnull:
        try:
            if subprocess.call(["btrfs", "subvolume", "list", path],
                               stdout=devnull, stderr=devnull) == 0:
                return "btrfs"
        except:
            pass

        if subprocess.call(["modinfo", "overlay"],
                           stdout=devnull, stderr=devnull) == 0:
            return "overlay"

        if subprocess.call(["modinfo", "overlayfs"],
                           stdout=devnull, stderr=devnull) == 0:
            return "overlayfs"

        return "directory"


def create_chroot(args):
    chroot = args.name

    print("[%s] Creating config" % chroot)
    config_path = os.path.join("/etc/schroot/chroot.d/", chroot)

    if os.path.exists(config_path):
        print("[%s] Chroot already exists!" % chroot)
        return

    host_arch = host_architecture()

    arch_config = {
        'amd64': {'i386': {'novirt': True, 'personality': 'linux32'}},
        'armhf': {},
        'arm64': {
            'armel': {'personality': 'linux32', 'novirt': True},
            'armhf': {'personality': 'linux32', 'novirt': True},
            },
        'i386': {},
        'powerpc': {},
        'ppc64el': {},
    }

    try:
        if arch_config[host_arch][args.architecture]['novirt']:
            pass
    except KeyError:
        os.system("sudo apt-get install qemu-user-static")

    personality = None
    try:
        personality = arch_config[host_arch][args.architecture]['personality']
    except KeyError:
        pass

    # Generate alias list
    aliases = []
    if args.series in args.name:
        for pocket in ("security", "updates", "proposed", "backports"):
            aliases.append(args.name.replace(args.series, "%s-%s" %
                           (args.series, pocket)))
            for component in ("main", "restricted", "universe", "multiverse"):
                aliases.append(args.name.replace(args.series, "%s-%s+%s" %
                               (args.series, pocket, component)))

    # Generate config
    config = ConfigParser()
    config.add_section(args.name)

    if args.backing_store == "auto":
        args.backing_store = guess_backingstore("/var/lib/schroot/")

    if args.backing_store in ("overlay", "overlayfs", "directory"):
        config.set(args.name, "type", "directory")
        config.set(args.name, "directory",
                   os.path.join("/var/lib/schroot/chroot/", args.name))

    if args.backing_store == "overlay":
        config.set(args.name, "union-type", "overlay")
    elif args.backing_store == "overlayfs":
        config.set(args.name, "union-type", "overlayfs")
    elif args.backing_store == "btrfs":
        snapshot_dir = os.path.join("/var/lib/schroot/snapshots/", args.name)
        if not os.path.exists(snapshot_dir):
            os.makedirs(snapshot_dir)

        config.set(args.name, "type", "btrfs-snapshot")
        config.set(args.name, "btrfs-source-subvolume",
                   os.path.join("/var/lib/schroot/chroot/", args.name))
        config.set(args.name, "btrfs-snapshot-directory",
                   os.path.join("/var/lib/schroot/snapshots/", args.name))

    config.set(args.name, "description", "Ubuntu %s/%s sbuild" %
               (args.series, args.architecture))
    config.set(args.name, "root-groups", "root,sbuild")
    config.set(args.name, "profile", "sbuild")
    config.set(args.name, "launchpad.dist", "ubuntu")
    config.set(args.name, "launchpad.series", args.series)
    config.set(args.name, "launchpad.arch", args.architecture)
    config.set(args.name, "launchpad.url", "")
    config.set(args.name, "apt.enable", "true")
    config.set(args.name, "aliases", ",".join(aliases))
    if personality:
        config.set(args.name, "personality", personality)

    with open(config_path, "w+") as fd:
        config.write(fd)

        fd.seek(0)
        content = fd.read()
        fd.truncate(0)
        fd.seek(0)
        fd.write(content.replace(" = ", "="))

    print("[%s] Initial download" % chroot)
    update_chroot(args)

    print("[%s] Done creating" % chroot)


def report_download(blocks_read, block_size, total_size):
    size_read = blocks_read * block_size
    percent = size_read/total_size*100
    print("%.0f %%" % (percent), end='\r')


def update_chroot(args):
    chroot = args.name

    # Lock file
    lock_file = "/run/lock/sbuild.%s" % chroot
    lock_fd = open(lock_file, 'w')
    try:
        fcntl.lockf(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except IOError:
        print("[%s] Waiting for other process to release the lock" % chroot)
        fcntl.lockf(lock_fd, fcntl.LOCK_EX)

    # Authenticated login to Launchpad
    from launchpadlib.launchpad import Launchpad
    lp = Launchpad.login_anonymously('lp-sbuild', 'production',
                                     launchpadlib_dir=cache_dir)

    print("[%s] Processing config" % chroot)
    config_path = os.path.join("/etc/schroot/chroot.d/", chroot)

    # Parse the existing configuration
    config = ConfigParser()
    config.read(config_path)

    if not os.path.exists(config_path):
        print("[%s] Doesn't exist." % chroot)
        return

    if (set(['launchpad.dist', 'launchpad.series', 'launchpad.arch'])
            - set(config.options(chroot))):
        print("[%s] Skipping (missing launchpad.* options)." % chroot)
        return

    # Grab the chroot URL from Launchpad
    dist = lp.distributions[config.get(chroot, 'launchpad.dist')]
    series = dist.getSeries(name_or_version=config.get(chroot,
                                                       'launchpad.series'))
    arch = series.getDistroArchSeries(archtag=config.get(chroot,
                                                         'launchpad.arch'))
    chroot_url = arch.chroot_url

    old_chroot_url = None
    if config.has_option(chroot, 'launchpad.url'):
        old_chroot_url = config.get(chroot, 'launchpad.url')

    if chroot_url == old_chroot_url:
        print("[%s] Already up to date." % chroot)
        return

    if not chroot_url.endswith(".tar.bz2"):
        print("[%s] Remote tarball isn't .tar.bz2." % chroot)
        return

    print("[%s] Downloading new Launchpad chroot." % chroot)

    # Create tempfile for chroot download
    fd, tarball_path = tempfile.mkstemp(suffix=".tar.bz2")
    os.remove(tarball_path)
    os.close(fd)

    # Download the chroot
    urllib.urlretrieve(chroot_url, tarball_path, report_download)

    # Unpack the chroot
    if config.has_option(chroot, "directory"):
        chroot_path = config.get(chroot, "directory")
    elif config.has_option(chroot, "btrfs-source-subvolume"):
        chroot_path = config.get(chroot, "btrfs-source-subvolume")

    if config.get(chroot, "type") == "btrfs-snapshot":
        if os.path.exists(chroot_path):
            with open(os.devnull, "w") as devnull:
                subprocess.call(["btrfs", "subvolume", "delete", chroot_path],
                                stdout=devnull, stderr=devnull)

        if not os.path.exists(os.path.dirname(chroot_path)):
            os.makedirs(os.path.dirname(chroot_path))

        with open(os.devnull, "w") as devnull:
            subprocess.call(["btrfs", "subvolume", "create", chroot_path],
                            stdout=devnull, stderr=devnull)
    else:
        if os.path.exists(chroot_path):
            shutil.rmtree(chroot_path)

        os.makedirs(chroot_path)

    subprocess.call(['tar', 'xf', tarball_path, '-C', chroot_path,
                     '--strip-components=1', 'chroot-autobuild'])
    os.remove(tarball_path)

    # Make the chroot sbuild compatible
    if os.path.exists(os.path.join(chroot_path, "CurrentlyBuilding")):
        os.remove(os.path.join(chroot_path, "CurrentlyBuilding"))

    if os.path.exists(os.path.join(chroot_path,
                                   "etc/apt/sources.list.d/bootstrap.list")):
        os.remove(os.path.join(chroot_path, "etc/apt/sources.list.d/"
                                            "bootstrap.list"))
    if not os.path.exists(os.path.join(chroot_path, "etc/schroot/setup.d")):
        os.makedirs(os.path.join(chroot_path, "etc/schroot/setup.d"))

    with open(os.path.join(chroot_path, "etc/apt/sources.list"), "r") as sources:
        lines = sources.readlines()

    old_series = lines[0].split(' ')[2]
    with open(os.path.join(chroot_path, "etc/apt/sources.list"), "w") as sources:
        for line in lines:
            sources.write(re.sub(old_series, series.name, line))

    # Update the symlink to the chroot (just in case)
    if not os.path.exists("/etc/sbuild/chroot"):
        os.makedirs("/etc/sbuild/chroot")
    symlink_path = os.path.join("/etc/sbuild/chroot/", chroot)
    if os.path.exists(symlink_path):
        os.remove(symlink_path)
    os.symlink(chroot_path, symlink_path)

    # Set the chroot URL and save the config
    config.set(chroot, 'launchpad.url', chroot_url)
    with open(config_path, "w+") as fd:
        config.write(fd)

        fd.seek(0)
        content = fd.read()
        fd.truncate(0)
        fd.seek(0)
        fd.write(content.replace(" = ", "="))

    print("[%s] Done updating from Launchpad" % chroot)

    # Remove lock file
    if os.path.exists(lock_file):
        os.remove(lock_file)


def remove_chroot(args):
    chroot = args.name

    print("[%s] Beginning removal" % chroot)
    config_path = os.path.join("/etc/schroot/chroot.d/", chroot)

    # Parse the existing configuration
    config = ConfigParser()
    config.read(config_path)

    if not os.path.exists(config_path):
        print("[%s] Doesn't exist." % chroot)
        return

    if (set(['launchpad.dist', 'launchpad.series', 'launchpad.arch'])
            - set(config.options(chroot))):
        print("[%s] Skipping (missing launchpad.* options)." % chroot)
        return

    # Actually remove it
    if config.has_option(chroot, "directory"):
        chroot_path = config.get(chroot, "directory")
    elif config.has_option(chroot, "btrfs-source-subvolume"):
        chroot_path = config.get(chroot, "btrfs-source-subvolume")

    if config.get(chroot, "type") == "btrfs-snapshot":
        if os.path.exists(chroot_path):
            with open(os.devnull, "w") as devnull:
                subprocess.call(["btrfs", "subvolume", "delete", chroot_path],
                                stdout=devnull, stderr=devnull)
        if os.path.exists(config.get(chroot, "btrfs-snapshot-directory")):
            shutil.rmtree(config.get(chroot, "btrfs-snapshot-directory"))
    else:
        shutil.rmtree(chroot_path)

    os.remove(config_path)

    symlink_path = os.path.join("/etc/sbuild/chroot/", chroot)
    if os.path.exists(symlink_path):
        os.remove(symlink_path)

    print("[%s] Removed" % chroot)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="launchpad sbuild chroot")
    subparsers = parser.add_subparsers()

    cache_dir = tempfile.mkdtemp()

    # Create
    parser_create = subparsers.add_parser(
        'create', help="Create a new Launchpad-based chroot")
    parser_create.add_argument("-n", "--name", required=True)
    parser_create.add_argument("-s", "--series", required=True)
    parser_create.add_argument("-a", "--architecture", required=True)
    parser_create.add_argument(
        "-b", "--backing-store", default="auto",
        choices=("auto", "btrfs", "directory", "overlay", "overlayfs"))

    parser_create.set_defaults(func=create_chroot)

    # Update
    parser_update = subparsers.add_parser(
        'update', help="Update a Launchpad-based chroot")
    parser_update.add_argument("-n", "--name", required=True)
    parser_update.set_defaults(func=update_chroot)

    # Remove
    parser_remove = subparsers.add_parser(
        'remove', help="Remove a Launchpad-based chroot")
    parser_remove.add_argument("-n", "--name", required=True)
    parser_remove.set_defaults(func=remove_chroot)

    args = parser.parse_args()

    # Some checks
    if not os.geteuid() == 0:
        parser.error("Sorry but you must be root.")

    args.func(args)

    if cache_dir and os.path.exists(cache_dir):
        shutil.rmtree(cache_dir)
