#!/usr/bin/env python

"""
This is a helper program of GToTree (https://github.com/AstrobioMike/GToTree/wiki)
for setting up reference files for the glorious Genome Taxonomy Database (gtdb.ecogenomic.org/).

It ensures the (slim) GTDB metadata table is present, downloading bit's
pre-slimmed tarball if it isn't. There is no version/freshness check: if the
table is present it is left alone. Rebuilding directly from the upstream GTDB
release happens only as a fallback if the hosted tarball can't be retrieved.

For examples, please visit the GToTree wiki here: https://github.com/AstrobioMike/GToTree/wiki/example-usage
"""

import sys
import os
import argparse
import importlib.util
from importlib.machinery import SourceFileLoader

_THIS_DIR = os.path.dirname(os.path.abspath(__file__))


def _load(module_basename):
    path = os.path.join(_THIS_DIR, module_basename)
    mod_name = module_basename.replace("-", "_")
    loader = SourceFileLoader(mod_name, path)
    spec = importlib.util.spec_from_loader(mod_name, loader)
    mod = importlib.util.module_from_spec(spec)
    loader.exec_module(mod)
    return mod


_gtdb = _load("gtt-gtdb-helpers")
_helpers = _load("gtt-ncbi-helpers")
color_text = _helpers.color_text
wprint = _helpers.wprint


parser = argparse.ArgumentParser(
    description="This is a helper program to facilitate setting up the reference files for the "
                "glorious Genome Taxonomy Database (gtdb.ecogenomic.org). It's really meant for "
                "internal use only by the main GToTree program.")

parser.add_argument("-f", "--force-update",
                    help="Re-download the slim GTDB table even if it is already present "
                         "(pulls the hosted slim tarball again).",
                    action="store_true")

args = parser.parse_args()

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

def main():

    ## checking env variable is set and writable
    GTDB_dir = check_location_var_is_set_and_writable("GTDB_dir")

    ## setting up ref GTDB files if needed (or if forced)
    if _gtdb.check_if_gtdb_data_present(GTDB_dir) and not args.force_update:
        sys.exit(0)

    _gtdb.get_slim_gtdb_tab(GTDB_dir)

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


def check_location_var_is_set_and_writable(variable):

    # making sure there is an env variable
    try:
        path = os.environ[variable]
        if path == "":
            raise KeyError
    except KeyError:
        print()
        wprint(color_text("The environment variable '" + str(variable) + "' does not seem to be set, and we need it if wanting to add GTDB taxonomic lineages :(", "red"))
        print()
        wprint("Try to set it with `gtt-data-locations set`, then run GToTree again.")
        print("\nExiting for now.\n")
        sys.exit(1)

    # making sure path is writable for the user
    if not os.access(path, os.W_OK):
        print()
        wprint(color_text("The environment variable '" + str(variable) + "' does not seem to be writable, and we need it to be if wanting to add GTDB taxonomic lineages :(", "red"))
        print()
        wprint("Try to set it somewhere else with `gtt-data-locations set`, then run GToTree again.")
        print("\nExiting for now.\n")
        sys.exit(1)

    return path


if __name__ == "__main__":
    main()
