#!/usr/bin/env python

"""
This is a helper program of GToTree (https://github.com/AstrobioMike/GToTree/wiki)
to download and set up the slim NCBI assembly-info table.

Default behavior: download the pre-built, pre-combined slim table (packaged as a
single .tar.gz and rebuilt weekly by bit's refresh GitHub Action, re-hosted at a
fixed release asset). This is fast and is what almost everyone gets.

`-f/--force-update` re-pulls that same slim tarball (it does NOT rebuild directly
from NCBI). A full rebuild from NCBI's GenBank + RefSeq summaries only happens as
a fallback, if the hosted tarball can't be retrieved or extracted.

There is no longer a 4-week freshness check: if the table is present it is left
alone unless `-f` is given.

Ex. usage: gtt-get-ncbi-assembly-tables
"""

import sys
import os
import socket
import shutil
import tarfile
import time
import urllib
import urllib.error
import argparse
from datetime import date

# vendored helpers + slim-builder live alongside this script in GToTree's bin/,
# which is on PATH; import them by adding this script's own dir to sys.path so
# the `-` in the filenames doesn't matter (import via importlib below).
import importlib.util
from importlib.machinery import SourceFileLoader

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


def _load(module_basename):
    # these sibling modules have no .py extension (they live in GToTree's bin/),
    # so give importlib an explicit source loader rather than relying on
    # extension-based detection (which returns no spec for extensionless files).
    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


_helpers = _load("gtt-ncbi-helpers")
_slim = _load("gtt-slim-ncbi-assembly-summary")

wprint = _helpers.wprint
color_text = _helpers.color_text
report_message = _helpers.report_message
notify_premature_exit = _helpers.notify_premature_exit
download_with_tqdm = _helpers.download_with_tqdm

build_slim_assembly_summary = _slim.build_slim_assembly_summary
write_date_retrieved = _slim.write_date_retrieved
GENBANK_URL = _slim.GENBANK_URL
REFSEQ_URL = _slim.REFSEQ_URL
TABLE_FILENAME = _slim.TABLE_FILENAME
DATE_FILENAME = _slim.DATE_FILENAME


# pre-slimmed, pre-combined NCBI assembly-info table + date-retrieved.txt,
# packaged as a single .tar.gz and rebuilt weekly by bit's
# refresh-ncbi-assembly-info GitHub Action, then re-hosted at this fixed asset
# URL. The archive holds exactly two files at its root:
#   ncbi-assembly-info.tsv   (slim, clean-header; see gtt-slim-ncbi-assembly-summary)
#   date-retrieved.txt       (build date, YYYY,MM,DD)
# GToTree depends on the bit-hosted asset (identical layout to what GToTree needs).
NCBI_ASSEMBLY_TARBALL_URL = "https://github.com/AstrobioMike/bit/releases/download/ncbi-assembly-info-latest/ncbi-assembly-info.tar.gz"


parser = argparse.ArgumentParser(
    description="This is a helper program to download and set up the slim NCBI "
                "assembly-info table if it is not present.",
    epilog="Ex. usage: gtt-get-ncbi-assembly-tables\n")

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

args = parser.parse_args()


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

def main():

    ncbi_dir = check_location_var_is_set()

    data_present = check_if_data_present(ncbi_dir)

    if data_present and not args.force_update:
        # present and not forced -> nothing to do
        sys.exit(0)

    get_slim_ncbi_assembly_data(ncbi_dir)


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


def check_location_var_is_set():

    try:
        ncbi_assembly_data_dir = os.environ['NCBI_assembly_data_dir']
    except KeyError:
        wprint(color_text("The environment variable 'NCBI_assembly_data_dir' does not seem to be set :(", "yellow"))
        wprint("This shouldn't happen, check on things with `gtt-data-locations check`.")
        print("")
        sys.exit(0)

    return ncbi_assembly_data_dir


def check_if_data_present(location):
    """
    True if both the slim table and date-retrieved.txt are present and non-empty.
    If either is missing/empty, any stray copy is cleaned up and we return False
    so a fresh copy is pulled.
    """
    table_path = os.path.join(str(location), TABLE_FILENAME)
    date_retrieved_path = os.path.join(str(location), DATE_FILENAME)

    def is_nonempty_file(p):
        return os.path.isfile(p) and os.path.getsize(p) > 0

    if not is_nonempty_file(table_path) or not is_nonempty_file(date_retrieved_path):
        for p in (table_path, date_retrieved_path):
            if os.path.exists(p) and os.path.isfile(p):
                os.remove(p)
        return False

    return True


def _download_with_retries(url, label, dest, attempts=4, retry_wait=3):
    """
    download `url` to `dest`, retrying up to `attempts` times on transient
    failures (timeouts, connection resets, transient errors), with a short wait
    between tries. A 404 is raised immediately (not retried). Raises the last
    error if all attempts fail.
    """
    last_err = None
    for attempt in range(1, attempts + 1):
        try:
            download_with_tqdm(url, label, dest)
            return
        except urllib.error.HTTPError as err:
            if err.code == 404:
                raise
            last_err = err
        except (urllib.error.URLError, socket.timeout, TimeoutError,
                ConnectionError, OSError) as err:
            last_err = err

        if attempt < attempts:
            wprint(color_text(
                f"    download failed (attempt {attempt}/{attempts}); retrying...",
                "yellow"))
            time.sleep(retry_wait)

    raise last_err


def get_slim_ncbi_assembly_data(location):
    """
    default path: download the pre-built slim NCBI assembly-info tarball and
    extract its two files (the slim table and date-retrieved.txt) into
    `location`. Falls back to rebuilding directly from NCBI
    (download_ncbi_assembly_summary_data) if no tarball URL is configured or the
    download/extract fails, so the user always ends up with a usable table.
    """
    table_path = os.path.join(location, TABLE_FILENAME)
    date_path = os.path.join(location, DATE_FILENAME)

    if not NCBI_ASSEMBLY_TARBALL_URL:
        download_ncbi_assembly_summary_data(location)
        return

    print(color_text("\n    Downloading the prepared NCBI assembly-info table (only needs to be done once)...\n", "yellow"))

    tarball_path = os.path.join(location, "ncbi-assembly-info.tar.gz")
    expected = {TABLE_FILENAME, DATE_FILENAME}

    default_timeout = socket.getdefaulttimeout()
    socket.setdefaulttimeout(30)
    try:
        _download_with_retries(
            NCBI_ASSEMBLY_TARBALL_URL, "        NCBI prepared data", tarball_path)

        # a truncated/corrupt download trips here (full-stream read via
        # getmembers), triggering the fallback below rather than writing a
        # corrupt table.
        with tarfile.open(tarball_path, "r:gz") as tar:
            members = tar.getmembers()
            # take only the two expected files, matched by basename, guarding
            # against path traversal / nested dirs by extracting to flat names.
            wanted = {}
            for m in members:
                base = os.path.basename(m.name)
                if base in expected and m.isfile():
                    wanted[base] = m
            missing = expected - set(wanted)
            if missing:
                raise ValueError(
                    f"prepared NCBI archive is missing expected file(s): {', '.join(sorted(missing))}")

            for base, m in wanted.items():
                src = tar.extractfile(m)
                if src is None:
                    raise ValueError(f"could not read '{base}' from prepared NCBI archive")
                dest = os.path.join(location, base)
                with src, open(dest, "wb") as out:
                    shutil.copyfileobj(src, out)

    except (urllib.error.URLError, socket.timeout, TimeoutError, ConnectionError,
            tarfile.TarError, ValueError, OSError) as err:
        # clean up partial artifacts, then fall back to rebuilding from NCBI
        for p in (tarball_path, table_path, date_path):
            if os.path.exists(p):
                try:
                    os.remove(p)
                except OSError:
                    pass
        print("")
        wprint(color_text("  Couldn't get the prepared NCBI assembly-info table; "
                          "rebuilding directly from NCBI instead.", "yellow"))
        report_message(f"Underlying issue: {err}", initial_indent="    ",
                       subsequent_indent="    ")
        print("")
        download_ncbi_assembly_summary_data(location)
        return
    finally:
        socket.setdefaulttimeout(default_timeout)
        if os.path.exists(tarball_path):
            try:
                os.remove(tarball_path)
            except OSError:
                pass
    print("")


def download_ncbi_assembly_summary_data(location):
    """
    rebuild fallback (used only if the hosted slim tarball can't be retrieved):
    download the GenBank + RefSeq assembly summaries directly from NCBI, combine
    and slim them to the columns GToTree uses (identical layout to the hosted
    asset), and stamp date-retrieved.txt with today's date.
    """
    genbank_temp = os.path.join(str(location), "assembly_summary_genbank.txt")
    refseq_temp = os.path.join(str(location), "assembly_summary_refseq.txt")
    table_path = os.path.join(str(location), TABLE_FILENAME)
    date_path = os.path.join(str(location), DATE_FILENAME)

    print(color_text("\n    Downloading NCBI assembly summaries (only needs to be done once)...\n", "yellow"))

    try:
        download_with_tqdm(GENBANK_URL, "        Genbank assemblies summary", genbank_temp)
        download_with_tqdm(REFSEQ_URL, "        RefSeq assemblies summary", refseq_temp)
        print("")
    except Exception as e:
        report_message(f"Downloading the NCBI assembly summary tables failed with the following error:\n{e}", "red")
        notify_premature_exit()
        return

    # combine + slim to the columns GToTree uses (clean-header layout)
    try:
        build_slim_assembly_summary(genbank_temp, refseq_temp, table_path)
    except Exception as e:
        report_message(f"Combining/slimming the NCBI assembly summary tables failed with the following error:\n{e}", "red")
        notify_premature_exit()
        return
    finally:
        for p in (genbank_temp, refseq_temp):
            if os.path.exists(p):
                os.remove(p)

    # storing date retrieved (YYYY,MM,DD), matching the asset's build stamp
    write_date_retrieved(date_path, when=date.today())


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

if __name__ == "__main__":
    main()
