#!/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/>.

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

from ConfigParser import ConfigParser


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

    # Generate alias list
    aliases = []
    for pocket in ("security", "updates", "proposed"):
        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)
    config.set(args.name, "type", "directory")
    config.set(args.name, "description", "Ubuntu %s/%s sbuild" %
               (args.series, args.architecture))
    config.set(args.name, "directory",
               os.path.join("/var/lib/schroot/chroot/", args.name))
    config.set(args.name, "union-type", "overlayfs")
    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))

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

    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)

    # Unpack the chroot
    chroot_path = config.get(chroot, 'directory')
    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"))

    # 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
    chroot_path = config.get(chroot, 'directory')
    if os.path.exists(chroot_path):
        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()

    # 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("-a", "--architecture", required=True)
    parser_create.add_argument("-s", "--series", required=True)
    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)
