#!/usr/bin/env python

# Name: whisk
# Purpose: A proof-of-concept "Working Environment" builder
# Author: Ken McIvor <mcivor@iit.edu>
#
# Copyright 2006 Illinois Institute of Technology
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL ILLINOIS INSTITUTE OF TECHNOLOGY BE LIABLE FOR ANY
# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#
# Except as contained in this notice, the name of Illinois Institute
# of Technology shall not be used in advertising or otherwise to promote
# the sale, use or other dealings in this Software without prior written
# authorization from Illinois Institute of Technology.

import sys
import os.path
import optparse
import tempfile
import distutils.log
import pkg_resources
import setuptools.command.easy_install

__version__ = '1.0'

WHISK_BOOTSTRAP = '''\
# BEGIN WHISK BOOTSTRAP # # # # # # # # # # # # # # # # # # # # # # # # # # # #
import sys
import os.path
if sys.path[0] != '':
    whisk_lib = os.path.join(os.path.dirname(sys.path[0]), 'lib')
    if os.path.isdir(whisk_lib):
        setuptools_pth = os.path.join(whisk_lib, %r)
        if os.path.exists(setuptools_pth):
            sys.path[0] = whisk_lib
            sys.path.insert(1, setuptools_pth)
        else:
            sys.exit('%%s: setuptools is missing from the whisk library'
                %% os.path.basename(sys.argv[0] or sys.executable))
    else:
        sys.exit('%%s: whisk library directory is missing'
            %% os.path.basename(sys.argv[0] or sys.executable))
else:
    sys.exit('%%s: this script cannot be run from stdin'
        %% os.path.basename(sys.argv[0] or sys.executable))
# END WHISK BOOTSTRAP # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
'''


def parse_arguments(argv):
    USAGE = '%prog [OPTION] ... REQUIREMENT_OR_FILE_OR_URL ...'
    VERSION = '%prog ' + __version__ + ', by Ken McIvor <mcivor@iit.edu>'
    parser = optparse.OptionParser(usage=USAGE, version=VERSION)

    parser.add_option('-d', '--install-dir', action='store',
        dest='install_dir', help='install environment in directory D',
        metavar='D')

    parser.add_option('-U', '--upgrade', action='store_true', dest='upgrade',
        help='force upgrade (searches PyPI for latest versions)')

    parser.add_option('-x', '--exclude-scripts', action='store_false',
        dest='scripts', default=True, help='don\'t install scripts')

    opts, args = parser.parse_args(argv)
    if not len(args):
        parser.print_usage()
        sys.exit(1)

    return opts, args


class EggWrangler:
    def __init__(self, upgrade, scripts):
        self.upgrade = upgrade
        self.scripts = scripts
        self.distributions = []
        self.tmpdir = tempfile.mkdtemp(prefix='whisk-')
        self.libdir, self.bindir = self._make_subdirs(self.tmpdir)

    def cleanup(self):
        self._rmtree(self.tmpdir)

    def _rmtree(self, path):
        if os.path.isdir(path):
            setuptools.command.easy_install.rmtree(path)

    def _make_subdirs(self, tmpdir):
        libdir = os.path.join(tmpdir, 'lib')
        bindir = os.path.join(tmpdir, 'bin')
        os.mkdir(libdir)
        os.mkdir(bindir)
        return libdir, bindir

    def _move_directory_contents(self, src, dst):
        for name in os.listdir(src):
            os.rename(os.path.join(src, name), os.path.join(dst, name))

    def _install_lib(self, src):
        # build the list of eggs, then deploy them
        egg_files = [os.path.basename(x.location)
                        for x in pkg_resources.find_distributions(src)]

        self._move_directory_contents(src, self.libdir)

        # activate the new distributions so setuptools knows they exist
        distributions = []
        for egg_file in egg_files:
            dist = pkg_resources.Distribution.from_filename(
                    os.path.join(self.libdir, egg_file))
            dist.activate()
            distributions.append(dist)

        self.distributions.extend(distributions)
        return distributions[0]

    def _install_bin(self, src):
        self._move_directory_contents(src, self.bindir)

    def _process_script(self, name, setuptools_dir):
        if name.endswith('.exe'):
            return

        # insert the WHISK_BOOTSTRAP code after the first line of the script
        script = file(name, 'r')
        contents = script.readlines()
        script.close()

        script = file(name, 'w')
        script.write(contents[0])
        script.write(WHISK_BOOTSTRAP % setuptools_dir)
        for line in contents[1:]:
            script.write(line)
        script.close()

    def easy_install(self, requirement):
        tmpdir = tempfile.mkdtemp(prefix='whisk-')
        libdir, bindir = self._make_subdirs(tmpdir)

        # hack setuptools to disable .pth checking
        easy_install_cls = setuptools.command.easy_install.easy_install
        old_check_site_dir = getattr(easy_install_cls, 'check_site_dir', None)
        easy_install_cls.check_site_dir = lambda *args, **kwds: True

        argv = ['--install-dir=%s' % libdir, '--always-copy']
        if self.upgrade:
            argv.append('--upgrade')

        if self.scripts:
            argv.append('--script-dir=%s' % bindir)
        else:
            argv.append('--exclude-scripts')
        argv.append(requirement)

        try:
            setuptools.command.easy_install.main(argv)
            self._install_bin(bindir)
            return self._install_lib(libdir)
        finally:
            # unhack setuptools and remove the tmpdir
            easy_install_cls.check_site_dir = old_check_site_dir
            self._rmtree(tmpdir)

    def _get_setuptools_dir(self):
        # return the directory under 'lib/' that setuptools was installed in
        setuptools_dist = None
        for dist in self.distributions:
            if dist.project_name == 'setuptools':
                setuptools_dist = dist

        # install setuptools if necessary
        if setuptools_dist is None:
            setuptools_dist = self.easy_install('setuptools')

        return os.path.basename(setuptools_dist.location)

    def install_env(self, install_dir):
        setuptools_dir = self._get_setuptools_dir()
        for name in os.listdir(self.bindir):
            self._process_script(os.path.join(self.bindir, name),
                setuptools_dir)

        if os.path.exists(install_dir):
            self._rmtree(install_dir)

        os.rename(self.tmpdir, install_dir)


def main(options, args):
    distutils.log.set_verbosity(1)
    wrangler = EggWrangler(options.upgrade, options.scripts)

    try:
        for target in args:
            wrangler.easy_install(target)

        if options.install_dir is None:
            # by default, the environment is named after the first distribution
            egg_name = wrangler.distributions[0].egg_name()
            if [x.platform for x in wrangler.distributions if x.platform]:
                egg_name += '-' + pkg_resources.get_platform()
            wrangler.install_env(egg_name)
        else:
            wrangler.install_env(options.install_dir)
    finally:
        wrangler.cleanup()


if __name__ == '__main__':
    try:
        options, args = parse_arguments(sys.argv[1:])
        main(options, args)
    except KeyboardInterrupt:
        pass