# Copyright (c) 2008-2011 by Enthought, Inc.
# All rights reserved.
#
# Author: Ilan Schnell <ischnell@enthought.com>
#
# This program unpacks/installs a single zip-file.  It also removes files
# that were installed, see below.

import os
import sys
import zipfile
from os.path import dirname, join, isdir


PREFIX = sys.prefix


def unpack(zip_path):
    """
    Unpack the zip file into PREFIX, creating directories as required,
    and writing the .txt file.
    """
    assert zip_path.endswith('.zip')
    txt_path = zip_path[:-4] + '.txt'

    z = zipfile.ZipFile(zip_path)
    fo_txt = open(txt_path, 'w')
    for name in z.namelist():
        if name.endswith('/') or name.startswith('.unused'):
            continue
        path = join(PREFIX, *name.split('/'))
        dir_path = dirname(path)
        if not isdir(dir_path):
            os.makedirs(dir_path)
        fo = open(path, 'wb')
        fo.write(z.read(name))
        fo.close()
        fo_txt.write("%s\n" % name)

    z.close()
    fo_txt.close()


def _unlink(path):
    try:
        os.unlink(path)
    except OSError:
        pass


def remove(txt_path):
    """
    Removes all files listed in FILE.txt.
    Empty directories are removed as required.
    Finally, removes txt_path itself.
    """
    assert txt_path.endswith('.txt')

    dirs1 = set()
    for line in open(txt_path):
        name = line.strip()
        if not name:
            continue
        path = join(PREFIX, *name.split('/'))
        _unlink(path)
        if path.endswith('.py'):
            _unlink(path + 'c')
        dirs1.add(dirname(path))

    dirs2 = set()
    for path in dirs1:
        while len(path) > len(PREFIX):
            dirs2.add(path)
            path = dirname(path)

    for path in sorted(dirs2, key=len, reverse=True):
        try:
            os.rmdir(path)
        except OSError: # directory might not exist or not be empty
            pass

    _unlink(txt_path)


def main():
    from optparse import OptionParser

    p = OptionParser(usage = "usage: %prog [options] FILE.zip|FILE.txt",
                     description =  """\
Enthought install utility for zip-files.  This program is called during the
install progress, after a zip-file has been written.

This program installs all files contained in FILE.zip into sys.prefix (or the
current working directory, option --cwd), creating directories as required.
A file named FILE.txt is created, which contains the relative path (always
using '/' as the separator, i.e. independent of the platform) for each file
written.

When the argument is FILE.txt, the program removes all files listed in
FILE.txt (removing empty directories as required), and finally it also
removes FILE.txt itself.
""")

    p.add_option("--cwd",
                 action="store_true",
                 help="unzip/removing in current working directory")

    opts, args = p.parse_args()

    if len(args) != 1:
        p.error("exactly one arument expected")

    if opts.cwd:
        global PREFIX
        PREFIX = os.getcwd()

    path = args[0]

    if path.endswith('.zip'):
        print "Extracting: %r" % path
        unpack(path)
    elif path.endswith('.txt'):
        print "Removing files listed in %r" % path
        remove(path)
    else:
        p.error(".zip or .txt file expected")


if __name__ == '__main__':
    main()