#!/usr/bin/env python3

# SPDX-FileCopyrightText: 2009 Fermi Research Alliance, LLC
# SPDX-License-Identifier: Apache-2.0

"""This program creates a condor tarball suitable to be used in the glideins"""


import os
import os.path
import shutil
import sys
import tarfile
import tempfile
import traceback

from glideinwms.creation.lib import cgWCreate

STARTUP_DIR = sys.path[0]
sys.path.append(os.path.join(STARTUP_DIR, "../.."))


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 Exception:
                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 Exception:
            usage()
            tb = traceback.format_exception(sys.exc_info()[0], sys.exc_info()[1], sys.exc_info()[2])
            print("\n".join(tb))
            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 Exception:
                    # 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)
