#!/usr/bin/env python

#
# Description:
#  This program creates a condor tarball
#  suitable to be used in the glideins
#

from __future__ import print_function
import os
import sys
import tempfile
import shutil
import tarfile
import os.path
import string
import traceback
STARTUP_DIR = sys.path[0]
sys.path.append(os.path.join(STARTUP_DIR, "../.."))
from glideinwms.creation.lib import cgWCreate


def usage():
    print("""
Usage: create_condor_tarball output_filename input_tarfile_or_dir

   output_filename - file name of the output tarball created

   input_tarfile_or_dir - this is the source tarfile or if already expanded,
                          the top level directory.

This utility can be used to create an condor tarball for use by the glideinwms
pilot process for various entry points.  The glideinwms pilot does not need all
the files/processes contained in a normal HTCondor tarfile.  This utility 
extracts only the needed processes/libraries thereby reducing the space needed 
on your factory node.

This takes an HTCondor tarfile (or expanded one) as input.
Extracts the needed binaries and creates a smaller tarball.
Refer to the glideinwms documentation:
   WMSfactory -> Configuration -> Multiple HTCondor Tarballs
for additional information on configuring the factory to use the tarball.
""")


################################################################################

def main(argv):
    if len(argv) != 3:
        usage()
        sys.exit(1)

    out_fname = argv[1]
    in_dir = argv[2]

    tmpdir = None
    if os.path.isfile(in_dir):
        # this should be a tarfile, not a dir
        # Untar in a tmpdir
        with tarfile.open(in_dir,"r:gz") as fd:
            tmpdir = tempfile.mkdtemp("_tmp","condor_",os.path.dirname(out_fname))
            try:
                # first create the regular files
                for f in fd.getmembers():
                    if not f.islnk():
                        fd.extract(f,tmpdir)
                # then create the links
                for f in fd.getmembers():
                    if f.islnk():
                        os.link(os.path.join(tmpdir, f.linkname), os.path.join(tmpdir, f.name))
                in_dir = os.path.join(tmpdir, os.listdir(tmpdir)[0])  # there is a condor* subdir there
            except:
                shutil.rmtree(tmpdir)
                raise
    elif os.path.isdir(in_dir):
        pass  # good...it is a directory
    else:
        print("%s does not exist!" % in_dir)
        sys.exit(2)

    try:
        try:
            tar_fd = cgWCreate.create_condor_tar_fd(in_dir)
        except RuntimeError as e:
            print("%s" % e)
            sys.exit(2)

        try:
            out_fd = open(out_fname, "w")
        except:
            usage()
            tb = traceback.format_exception(sys.exc_info()[0], sys.exc_info()[1],
                                            sys.exc_info()[2])
            print(string.join(tb, '\n'))
            sys.exit(3)

        delete_at_end = True  # cleanup by default
        try:
            out_fd.write(tar_fd.read())
            delete_at_end = False  # successful, preserve
        finally:
            out_fd.close()

            if delete_at_end:
                try:
                    os.unlink(out_fname)
                    print("Tarball %s deleted" % out_fname)
                except:
                    # never throw an error in the file deletion
                    pass 
    finally:
        if tmpdir is not None:
            # always remove the tmpdir before exiting
            shutil.rmtree(tmpdir)


############################################################
#
# S T A R T U P
# 
############################################################

if __name__ == '__main__':
    main(sys.argv)
